diff --git a/backend/data_export/tests/test_task.py b/backend/data_export/tests/test_task.py index 225ed06f..674a56f0 100644 --- a/backend/data_export/tests/test_task.py +++ b/backend/data_export/tests/test_task.py @@ -12,6 +12,7 @@ from projects.models import ( DOCUMENT_CLASSIFICATION, IMAGE_CLASSIFICATION, INTENT_DETECTION_AND_SLOT_FILLING, + SEGMENTATION, SEQ2SEQ, SEQUENCE_LABELING, SPEECH2TEXT, @@ -550,6 +551,77 @@ class TestExportBoundingBox(TestExport): self.assertEqual(dataset, expected_dataset) +class TestExportSegmentation(TestExport): + def prepare_data(self, collaborative=False): + self.project = prepare_project(SEGMENTATION, collaborative_annotation=collaborative) + self.example1 = mommy.make("ExportedExample", project=self.project.item, text="confirmed") + self.example2 = mommy.make("ExportedExample", project=self.project.item, text="unconfirmed") + self.seg1 = mommy.make("ExportedSegmentation", example=self.example1, user=self.project.admin, points=[0, 1]) + self.seg2 = mommy.make( + "ExportedSegmentation", example=self.example1, user=self.project.annotator, points=[1, 2] + ) + mommy.make("ExampleState", example=self.example1, confirmed_by=self.project.admin) + self.data1 = self.data_to_filename(self.example1) + self.data2 = self.data_to_filename(self.example2) + self.column = "segmentation" + + def test_unconfirmed_and_non_collaborative(self): + self.prepare_data() + datasets = self.export_dataset() + expected_datasets = { + self.project.admin.username: [ + { + **self.data1, + self.column: [self.seg1.to_dict()], + }, + {**self.data2, self.column: []}, + ], + self.project.approver.username: [ + {**self.data1, self.column: []}, + {**self.data2, self.column: []}, + ], + self.project.annotator.username: [ + { + **self.data1, + self.column: [self.seg2.to_dict()], + }, + {**self.data2, self.column: []}, + ], + } + for username, dataset in expected_datasets.items(): + self.assertEqual(datasets[username], dataset) + + def test_unconfirmed_and_collaborative(self): + self.prepare_data(collaborative=True) + dataset = self.export_dataset() + expected_dataset = [ + { + **self.data1, + self.column: [self.seg1.to_dict(), self.seg2.to_dict()], + }, + {**self.data2, self.column: []}, + ] + self.assertEqual(dataset, expected_dataset) + + def test_confirmed_and_non_collaborative(self): + self.prepare_data() + datasets = self.export_dataset(confirmed_only=True) + expected_datasets = {self.project.admin.username: [{**self.data1, self.column: [self.seg1.to_dict()]}]} + for username, dataset in expected_datasets.items(): + self.assertEqual(datasets[username], dataset) + + def test_confirmed_and_collaborative(self): + self.prepare_data(collaborative=True) + dataset = self.export_dataset(confirmed_only=True) + expected_dataset = [ + { + **self.data1, + self.column: [self.seg1.to_dict(), self.seg2.to_dict()], + } + ] + self.assertEqual(dataset, expected_dataset) + + class TestExportRelation(TestExport): def prepare_data(self, collaborative=False): self.project = prepare_project(SEQUENCE_LABELING, use_relation=True, collaborative_annotation=collaborative)