Serializers: operations/context records - filter vy context records

3 files changed, 102 insertions, 16 deletions

diff --git a/archaeological_context_records/serializers.py b/archaeological_context_records/serializers.py
index 98f7d2226..6035f0182 100644
--- a/archaeological_context_records/serializers.py
+++ b/archaeological_context_records/serializers.py
@@ -13,24 +13,40 @@ def cr_serialization(archive=False, return_empty_types=False,
                             archive_name=None, operation_queryset=None,
                             site_queryset=None, cr_queryset=None):
     result_queryset = {}
-    """
     if operation_queryset:
+        operation_ids = operation_queryset.values_list("id", flat=True)
         result_queryset = {
-            models.ContextRecord.__name__: operation_queryset,
-            models.ArchaeologicalSite.__name__:
-                models.ArchaeologicalSite.objects.filter(
-                    operations__id__in=operation_queryset.values_list(
-                        "id", flat=True)
-                    )
+            models.ContextRecord.__name__:
+                models.ContextRecord.objects.filter(
+                    operation_id__in=operation_ids),
+            models.RecordRelations.__name__:
+                models.RecordRelations.objects.filter(
+                    left_record__operation_id__in=operation_ids,
+                    right_record__operation_id__in=operation_ids,
+                )
         }
     elif site_queryset:
+        site_ids = site_queryset.values_list("id", flat=True)
         result_queryset = {
-            models.ArchaeologicalSite.__name__: site_queryset,
-            models.Operation.__name__: models.Operation.objects.filter(
-                archaeological_sites__id__in=site_queryset.values_list(
-                    "id", flat=True))
+            models.ContextRecord.__name__:
+                models.ContextRecord.objects.filter(
+                    archaeological_site_id__in=site_ids),
+            models.RecordRelations.__name__:
+                models.RecordRelations.objects.filter(
+                    left_record__archaeological_site_id__in=site_ids,
+                    right_record__archaeological_site_id__in=site_ids,
+                )
+        }
+    elif cr_queryset:
+        cr_ids = cr_queryset.values_list("id", flat=True)
+        result_queryset = {
+            models.ContextRecord.__name__: cr_queryset,
+            models.RecordRelations.__name__:
+                models.RecordRelations.objects.filter(
+                    left_record_id__in=cr_ids,
+                    right_record_id__in=cr_ids,
+                )
         }
-    """
 
     result = generic_get_results(CR_MODEL_LIST, "context_records",
                                  result_queryset=result_queryset)
diff --git a/archaeological_context_records/tests.py b/archaeological_context_records/tests.py
index 024a8f338..92f97545e 100644
--- a/archaeological_context_records/tests.py
+++ b/archaeological_context_records/tests.py
@@ -45,6 +45,8 @@ from ishtar_common.tests import WizardTest, WizardTestFormData as FormData, \
     OPERATION_TOWNS_FIXTURES, GenericSerializationTest, COMMON_FIXTURES, \
     WAREHOUSE_FIXTURES
 
+from archaeological_operations.models import Operation
+from archaeological_operations.serializers import operation_serialization
 from archaeological_context_records import views, serializers
 
 
@@ -210,8 +212,14 @@ class SerializationTest(GenericSerializationTest, ContextRecordInit, TestCase):
     fixtures = COMMON_FIXTURES + WAREHOUSE_FIXTURES
 
     def setUp(self):
-        cr = self.create_context_record(data={"label": u"CR 1"})[0]
-        cr2 = self.create_context_record(data={"label": u"CR 2"})[1]
+        ope1 = self.create_operation()[0]
+        ope2 = self.create_operation()[1]
+        cr = self.create_context_record(
+            data={"label": u"CR 1", "operation": ope1}
+        )[0]
+        cr2 = self.create_context_record(
+            data={"label": u"CR 2", "operation": ope2}
+        )[1]
         dating = models.Dating.objects.create(
             period=models.Period.objects.all()[0],
         )
@@ -228,7 +236,55 @@ class SerializationTest(GenericSerializationTest, ContextRecordInit, TestCase):
         )
 
     def test_serialization(self):
-        self.generic_serialization_test(serializers.cr_serialization)
+        res = self.generic_serialization_test(serializers.cr_serialization)
+        cr_json = json.loads(
+            res[('context_records',
+                 'archaeological_context_records__ContextRecord')]
+        )
+        self.assertEqual(len(cr_json), 2)
+
+        result_queryset = Operation.objects.filter(uuid=self.operations[0].uuid)
+        res = self.generic_serialization_test(
+            serializers.cr_serialization, no_test=True,
+            kwargs={"operation_queryset": result_queryset}
+        )
+        cr_json = json.loads(
+            res[('context_records',
+                 'archaeological_context_records__ContextRecord')]
+        )
+        self.assertEqual(len(cr_json), 1)
+
+        result_queryset = models.ContextRecord.objects.filter(
+            uuid=self.context_records[0].uuid)
+        res = self.generic_serialization_test(
+            serializers.cr_serialization, no_test=True,
+            kwargs={"cr_queryset": result_queryset}
+        )
+        cr_json = json.loads(
+            res[('context_records',
+                 'archaeological_context_records__ContextRecord')]
+        )
+        self.assertEqual(len(cr_json), 1)
+
+    def test_ope_serialization_with_cr_filter(self):
+        res = self.generic_serialization_test(
+            operation_serialization, no_test=True,
+        )
+        ope_json = json.loads(
+            res[('operations', 'archaeological_operations__Operation')]
+        )
+        self.assertEqual(len(ope_json), 2)
+
+        result_queryset = models.ContextRecord.objects.filter(
+            uuid=self.context_records[0].uuid)
+        res = self.generic_serialization_test(
+            operation_serialization, no_test=True,
+            kwargs={"cr_queryset": result_queryset}
+        )
+        ope_json = json.loads(
+            res[('operations', 'archaeological_operations__Operation')]
+        )
+        self.assertEqual(len(ope_json), 1)
 
     def test_restore(self):
         current_number, zip_filename = self.generic_restore_test_genzip(
diff --git a/archaeological_operations/serializers.py b/archaeological_operations/serializers.py
index 3500cdad2..a480f307a 100644
--- a/archaeological_operations/serializers.py
+++ b/archaeological_operations/serializers.py
@@ -12,7 +12,7 @@ OPERATION_MODEL_LIST = [
 
 def operation_serialization(archive=False, return_empty_types=False,
                             archive_name=None, operation_queryset=None,
-                            site_queryset=None):
+                            site_queryset=None, cr_queryset=None):
     result_queryset = {}
     if operation_queryset:
         operation_ids = operation_queryset.values_list("id", flat=True)
@@ -40,6 +40,20 @@ def operation_serialization(archive=False, return_empty_types=False,
                     right_record__archaeological_sites__id__in=site_ids,
                 ),
         }
+    elif cr_queryset:
+        cr_ids = cr_queryset.values_list("id", flat=True)
+        result_queryset = {
+            models.ArchaeologicalSite.__name__:
+                models.ArchaeologicalSite.objects.filter(
+                    operations__context_record__id__in=cr_ids),
+            models.Operation.__name__: models.Operation.objects.filter(
+                context_record__id__in=cr_ids),
+            models.RecordRelations.__name__:
+                models.RecordRelations.objects.filter(
+                    left_record__context_record__id__in=cr_ids,
+                    right_record__context_record__id__in=cr_ids,
+                ),
+        }
 
     result = generic_get_results(OPERATION_MODEL_LIST, "operations",
                                  result_queryset=result_queryset)