From 0eff4f62dee7e96209e4ba15da2c0e2e6bf58a72 Mon Sep 17 00:00:00 2001
From: Twin Karmakharm <hi@twin.uk.com>
Date: Fri, 26 May 2023 16:59:49 +0100
Subject: [PATCH] Fixed implementation of export tests

---
 backend/tests/test_models.py | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/backend/tests/test_models.py b/backend/tests/test_models.py
index 7821f065..384e4b64 100644
--- a/backend/tests/test_models.py
+++ b/backend/tests/test_models.py
@@ -1121,7 +1121,7 @@ def setUp(self):
                         "existing_annotator1": {
                             "sentiment": "positive"
                         },
-                        f"2": {
+                        f"{self.annotators[0].pk}": {
                             "sentiment": "positive"
                         }
 
@@ -1142,8 +1142,8 @@ def setUp(self):
                             ],
                             "next_annid": 1
                         },
-                        f"{settings.ANONYMIZATION_PREFIX}2": {
-                            "name": f"{settings.ANONYMIZATION_PREFIX}1",
+                        f"{settings.ANONYMIZATION_PREFIX}{self.annotators[0].pk}": {
+                            "name": f"{settings.ANONYMIZATION_PREFIX}{self.annotators[0].pk}",
                             "annotations": [
                                 {
                                     "type": "Document",
@@ -1250,9 +1250,9 @@ def test_export_gate_with_no_offset_type(self):
             self.assertEqual(doc_dict["offset_type"], "p", "offset_type should default to p")
 
 
-    def check_raw_gate_annotation_formatting(self, doc_dict):
+    def check_raw_gate_annotation_formatting(self, doc_dict: dict):
         self.assertTrue("annotation_sets" in doc_dict)
-        self.assertEqual(len(doc_dict["annotation_sets"]), 4)
+        self.assertEqual(len(doc_dict["annotation_sets"]), 4, doc_dict)
 
         # Test annotation formatting
         for aset_key, aset_data in doc_dict["annotation_sets"].items():
@@ -1288,7 +1288,7 @@ def test_export_csv(self):
             self.assertTrue("feature2" in doc_dict)
             self.assertTrue("feature3" in doc_dict)
             self.assertTrue("annotations" in doc_dict)
-            self.assertEqual(len(doc_dict["annotations"]), 4)
+            self.assertEqual(len(doc_dict["annotations"]), 4, doc_dict)
             anno_set_dict = doc_dict["annotations"]
             for set_key in anno_set_dict:
                 if set_key != "existing_annotator1":
@@ -1300,6 +1300,10 @@ def test_export_csv(self):
     def test_export_raw_anonymized(self):
 
         for document in self.project.documents.all():
+            # Mask any existing annotations that came with the document upload
+            document.data.pop("annotation_sets")
+            document.save()
+
             doc_dict = document.get_doc_annotation_dict("raw", anonymize=True)
 
             for aset_key, aset_data in doc_dict["annotation_sets"].items():
@@ -1309,6 +1313,10 @@ def test_export_raw_anonymized(self):
     def test_export_raw_deanonymized(self):
 
         for document in self.project.documents.all():
+            # Mask any existing annotations that came with the document upload
+            document.data.pop("annotation_sets")
+            document.save()
+
             doc_dict = document.get_doc_annotation_dict("raw", anonymize=False)
 
             for aset_key, aset_data in doc_dict["annotation_sets"].items():
@@ -1318,6 +1326,10 @@ def test_export_raw_deanonymized(self):
     def test_export_gate_anonymized(self):
 
         for document in self.project.documents.all():
+            # Mask any existing annotations that came with the document upload
+            document.data.pop("annotation_sets")
+            document.save()
+
             doc_dict = document.get_doc_annotation_dict("gate", anonymize=True)
 
             for aset_key, aset_data in doc_dict["annotation_sets"].items():
@@ -1327,6 +1339,10 @@ def test_export_gate_anonymized(self):
     def test_export_gate_deanonymized(self):
 
         for document in self.project.documents.all():
+            # Mask any existing annotations that came with the document upload
+            document.data.pop("annotation_sets")
+            document.save()
+
             doc_dict = document.get_doc_annotation_dict("gate", anonymize=False)
 
             for aset_key, aset_data in doc_dict["annotation_sets"].items():