Add test cases for data export of segmentation

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)