Serialization: warehouse - manage filter for context record, site, find and warehouse query

2 files changed, 130 insertions, 81 deletions

diff --git a/archaeological_warehouse/serializers.py b/archaeological_warehouse/serializers.py
index 9a49d5acf..2c44c9dcf 100644
--- a/archaeological_warehouse/serializers.py
+++ b/archaeological_warehouse/serializers.py
@@ -13,6 +13,48 @@ WAREHOUSE_MODEL_LIST = [
 # TODO: associated documents
 
 
+def generate_queryset(base_query_key, ids):
+    warehouse_division_q, warehouse_q, container_q = None, None, None
+    container_loca_q = None
+    for find_key in ("finds", "finds_ref"):
+        base_key = "{}__{}".format(find_key, base_query_key)
+        q = Q(**{base_key: ids})
+        q_loca = Q(
+            **{"container__{}__{}".format(find_key, base_query_key): ids})
+        if not container_q:
+            container_q = q
+            container_loca_q = q_loca
+        else:
+            container_q |= q
+            container_loca_q |= q_loca
+        for container_key in ("containers", "owned_containers"):
+            q = Q(**
+                  {"{}__{}__{}".format(container_key, find_key,
+                                       base_query_key): ids})
+            q_div = Q(**
+                      {"warehouse__{}__{}__{}".format(
+                          container_key, find_key, base_query_key): ids})
+            if not warehouse_q:
+                warehouse_q = q
+                warehouse_division_q = q_div
+            else:
+                warehouse_q |= q
+                warehouse_division_q |= q_div
+    result_queryset = {
+        models.Warehouse.__name__:
+            models.Warehouse.objects.filter(warehouse_q),
+        models.Container.__name__:
+            models.Container.objects.filter(container_q),
+        models.WarehouseDivisionLink.__name__:
+            models.WarehouseDivisionLink.objects.filter(
+                warehouse_division_q),
+        models.ContainerLocalisation.__name__:
+            models.ContainerLocalisation.objects.filter(
+                container_loca_q)
+    }
+    return result_queryset
+
+
 def warehouse_serialization(archive=False, return_empty_types=False,
                             archive_name=None, operation_queryset=None,
                             site_queryset=None, cr_queryset=None,
@@ -21,74 +63,35 @@ def warehouse_serialization(archive=False, return_empty_types=False,
     if operation_queryset:
         operation_ids = operation_queryset.values_list("id", flat=True)
         base_query_key = "base_finds__context_record__operation_id__in"
-        warehouse_division_q, warehouse_q, container_q = None, None, None
-        container_loca_q = None
-        for find_key in ("finds", "finds_ref"):
-            q = Q(**{"{}__{}".format(find_key, base_query_key): operation_ids})
-            q_loca = Q(
-                **{"container__{}__{}".format(find_key, base_query_key):
-                   operation_ids})
-            if not container_q:
-                container_q = q
-                container_loca_q = q_loca
-            else:
-                container_q |= q
-                container_loca_q |= q_loca
-            for container_key in ("containers", "owned_containers"):
-                q = Q(**
-                      {"{}__{}__{}".format(container_key, find_key,
-                                           base_query_key): operation_ids})
-                q_div = Q(**
-                          {"warehouse__{}__{}__{}".format(
-                              container_key, find_key, base_query_key):
-                           operation_ids})
-                if not warehouse_q:
-                    warehouse_q = q
-                    warehouse_division_q = q_div
-                else:
-                    warehouse_q |= q
-                    warehouse_division_q |= q_div
-        result_queryset = {
-            models.Warehouse.__name__:
-                models.Warehouse.objects.filter(warehouse_q),
-            models.Container.__name__:
-                models.Container.objects.filter(container_q),
-            models.WarehouseDivisionLink.__name__:
-                models.WarehouseDivisionLink.objects.filter(
-                    warehouse_division_q),
-            models.ContainerLocalisation.__name__:
-                models.ContainerLocalisation.objects.filter(
-                    container_loca_q)
-        }
+        result_queryset = generate_queryset(base_query_key, operation_ids)
     elif site_queryset:
-        sites = site_queryset.values_list("id", flat=True)
-        f_q = {
-            "base_finds__context_record__operation__archaeological_site_id__in": sites
-        }
-        result_queryset = {
-            models.BaseFind.__name__:
-                models.BaseFind.objects.filter(
-                    context_record__operation__archaeological_site_id__in=sites
-                ),
-            models.Find.__name__:
-                models.Find.objects.filter(**f_q),
-        }
+        site_ids = site_queryset.values_list("id", flat=True)
+        base_query_key = "base_finds__context_record__"\
+                         "archaeological_site_id__in"
+        result_queryset = generate_queryset(base_query_key, site_ids)
     elif cr_queryset:
         cr_ids = cr_queryset.values_list("id", flat=True)
-        result_queryset = {
-            models.BaseFind.__name__:
-                models.BaseFind.objects.filter(
-                    context_record__in=cr_ids),
-            models.Find.__name__:
-                models.Find.objects.filter(
-                    base_finds__context_record__in=cr_ids),
-        }
+        base_query_key = "base_finds__context_record__id__in"
+        result_queryset = generate_queryset(base_query_key, cr_ids)
     elif find_queryset:
         find_ids = find_queryset.values_list("id", flat=True)
+        base_query_key = "id__in"
+        result_queryset = generate_queryset(base_query_key, find_ids)
+    elif warehouse_queryset:
+        warehouse_ids = warehouse_queryset.values_list("id", flat=True)
         result_queryset = {
-            models.BaseFind.__name__: models.BaseFind.objects.filter(
-                find__id__in=find_ids),
-            models.Find.__name__: find_queryset,
+            models.Warehouse.__name__: warehouse_queryset,
+            models.Container.__name__:
+                models.Container.objects.filter(
+                    Q(location__id__in=warehouse_ids) |
+                    Q(responsible__id__in=warehouse_ids)
+                ),
+            models.WarehouseDivisionLink.__name__:
+                models.WarehouseDivisionLink.objects.filter(
+                    warehouse_id__in=warehouse_ids),
+            models.ContainerLocalisation.__name__:
+                models.ContainerLocalisation.objects.filter(
+                    division__warehouse_id__in=warehouse_ids)
         }
     result = generic_get_results(WAREHOUSE_MODEL_LIST, "warehouse",
                                  result_queryset=result_queryset)
diff --git a/archaeological_warehouse/tests.py b/archaeological_warehouse/tests.py
index a14c0fb8c..52f1a96fc 100644
--- a/archaeological_warehouse/tests.py
+++ b/archaeological_warehouse/tests.py
@@ -27,9 +27,12 @@ from ishtar_common.tests import WizardTest, WizardTestFormData as FormData, \
 from ishtar_common.models import IshtarSiteProfile, SpatialReferenceSystem
 from archaeological_operations.models import Operation
 from archaeological_context_records.models import ContextRecord
+from archaeological_finds.models import Find
 
 from archaeological_warehouse import models, views, forms, serializers
 
+from archaeological_operations.serializers import operation_serialization
+
 
 class SerializationTest(GenericSerializationTest, FindInit, TestCase):
     fixtures = COMMON_FIXTURES + WAREHOUSE_FIXTURES
@@ -61,6 +64,7 @@ class SerializationTest(GenericSerializationTest, FindInit, TestCase):
             external_id="test3",
             warehouse_type=models.WarehouseType.objects.all()[0],
         )
+        self.warehouses = [w1, w2, w3]
         c1 = models.Container.objects.create(
             location=w1,
             responsible=w1,
@@ -201,41 +205,83 @@ class SerializationTest(GenericSerializationTest, FindInit, TestCase):
                  "archaeological_warehouse__ContainerLocalisation")]
         )
         self.assertEqual(len(loca_json), 3)
-        """
 
         result_queryset = ContextRecord.objects.filter(
             uuid=self.context_records[0].uuid)
         res = self.generic_serialization_test(
-            serializers.find_serialization, no_test=True,
+            serializers.warehouse_serialization, no_test=True,
             kwargs={"cr_queryset": result_queryset}
         )
-        find_json = json.loads(
-            res[('finds', 'archaeological_finds__Find')]
+        warehouse_json = json.loads(
+            res[('warehouse', "archaeological_warehouse__Warehouse")]
+        )
+        self.assertEqual(len(warehouse_json), 2)
+        container_json = json.loads(
+            res[('warehouse', "archaeological_warehouse__Container")]
+        )
+        self.assertEqual(len(container_json), 2)
+        div_json = json.loads(
+            res[('warehouse',
+                 "archaeological_warehouse__WarehouseDivisionLink")]
         )
-        self.assertEqual(len(find_json), 1)
-        bfind_json = json.loads(
-            res[('finds', 'archaeological_finds__BaseFind')]
+        self.assertEqual(len(div_json), 3)
+        loca_json = json.loads(
+            res[('warehouse',
+                 "archaeological_warehouse__ContainerLocalisation")]
         )
-        self.assertEqual(len(bfind_json), 1)
+        self.assertEqual(len(loca_json), 3)
 
-        result_queryset = models.Find.objects.filter(
-            uuid=self.finds[0].uuid)
+        result_queryset = Find.objects.filter(uuid=self.finds[0].uuid)
         res = self.generic_serialization_test(
-            serializers.find_serialization, no_test=True,
+            serializers.warehouse_serialization, no_test=True,
             kwargs={"find_queryset": result_queryset}
         )
-        find_json = json.loads(
-            res[('finds', 'archaeological_finds__Find')]
+        warehouse_json = json.loads(
+            res[('warehouse', "archaeological_warehouse__Warehouse")]
         )
-        self.assertEqual(len(find_json), 1)
-        bfind_json = json.loads(
-            res[('finds', 'archaeological_finds__BaseFind')]
+        self.assertEqual(len(warehouse_json), 2)
+        container_json = json.loads(
+            res[('warehouse', "archaeological_warehouse__Container")]
+        )
+        self.assertEqual(len(container_json), 2)
+        div_json = json.loads(
+            res[('warehouse',
+                 "archaeological_warehouse__WarehouseDivisionLink")]
+        )
+        self.assertEqual(len(div_json), 3)
+        loca_json = json.loads(
+            res[('warehouse',
+                 "archaeological_warehouse__ContainerLocalisation")]
+        )
+        self.assertEqual(len(loca_json), 3)
+
+        result_queryset = models.Warehouse.objects.filter(
+            id=self.warehouses[0].id)
+        res = self.generic_serialization_test(
+            serializers.warehouse_serialization, no_test=True,
+            kwargs={"warehouse_queryset": result_queryset}
+        )
+        warehouse_json = json.loads(
+            res[('warehouse', "archaeological_warehouse__Warehouse")]
+        )
+        self.assertEqual(len(warehouse_json), 1)
+        container_json = json.loads(
+            res[('warehouse', "archaeological_warehouse__Container")]
+        )
+        self.assertEqual(len(container_json), 3)
+        div_json = json.loads(
+            res[('warehouse',
+                 "archaeological_warehouse__WarehouseDivisionLink")]
+        )
+        self.assertEqual(len(div_json), 2)
+        loca_json = json.loads(
+            res[('warehouse',
+                 "archaeological_warehouse__ContainerLocalisation")]
         )
-        self.assertEqual(len(bfind_json), 1)
-        """
+        self.assertEqual(len(loca_json), 2)
 
     """
-    def test_ope_serialization_with_find_filter(self):
+    def test_ope_serialization_with_warehouse_filter(self):
         res = self.generic_serialization_test(
             operation_serialization, no_test=True,
         )