Serializer: manage queryset filter

2 files changed, 62 insertions, 4 deletions

diff --git a/archaeological_operations/serializers.py b/archaeological_operations/serializers.py
index 7d4d9033d..22f0aedc7 100644
--- a/archaeological_operations/serializers.py
+++ b/archaeological_operations/serializers.py
@@ -9,9 +9,30 @@ OPERATION_MODEL_LIST = [
 
 
 def operation_serialization(archive=False, return_empty_types=False,
-                            archive_name=None):
-    result = generic_get_results(OPERATION_MODEL_LIST, "operations")
+                            archive_name=None, operation_queryset=None,
+                            site_queryset=None):
+    result_queryset = {}
+    if operation_queryset:
+        result_queryset = {
+            models.Operation.__name__: operation_queryset,
+            models.ArchaeologicalSite.__name__:
+                models.ArchaeologicalSite.objects.filter(
+                    operations__id__in=operation_queryset.values_list(
+                        "id", flat=True)
+                    )
+        }
+    elif site_queryset:
+        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))
+        }
+
+    result = generic_get_results(OPERATION_MODEL_LIST, "operations",
+                                 result_queryset=result_queryset)
     full_archive = archive_serialization(
         result, archive_dir="operations", archive=archive,
-        return_empty_types=return_empty_types, archive_name=archive_name)
+        return_empty_types=return_empty_types, archive_name=archive_name,
+    )
     return full_archive
diff --git a/archaeological_operations/tests.py b/archaeological_operations/tests.py
index d0726fbdc..2c06f5a35 100644
--- a/archaeological_operations/tests.py
+++ b/archaeological_operations/tests.py
@@ -750,9 +750,11 @@ class SerializationTest(GenericSerializationTest, TestCase):
         self.username, self.password, self.user = create_superuser()
         operation = create_operation(self.user,
                                      values={"code_patriarche": "66666"})
+        create_operation(self.user, values={"code_patriarche": "66667"})
         site = models.ArchaeologicalSite.objects.create(
             reference="ref-site"
         )
+        operation.archaeological_sites.add(site)
         operation.top_sites.add(site)
         default = {"town": Town.objects.create(numero_insee="66666"),
                    "section": 'A', 'parcel_number': '1', "operation": operation}
@@ -762,7 +764,42 @@ class SerializationTest(GenericSerializationTest, TestCase):
             start_date=datetime.date.today(), end_date=datetime.date.today())
 
     def test_serialization(self):
-        self.generic_serialization_test(serializers.operation_serialization)
+        res = self.generic_serialization_test(
+            serializers.operation_serialization)
+        ope_json = json.loads(
+            res[('operations', 'archaeological_operations__Operation')]
+        )
+        self.assertEqual(len(ope_json), 2)
+
+        result_queryset = models.Operation.objects.filter(
+            code_patriarche="66666")
+        res = self.generic_serialization_test(
+            serializers.operation_serialization, no_test=True,
+            kwargs={"operation_queryset": result_queryset}
+        )
+        ope_json = json.loads(
+            res[('operations', 'archaeological_operations__Operation')]
+        )
+        self.assertEqual(len(ope_json), 1)
+        site_json = json.loads(
+            res[('operations', 'archaeological_operations__ArchaeologicalSite')]
+        )
+        self.assertEqual(len(site_json), 1)
+
+        result_queryset = models.Operation.objects.filter(
+            code_patriarche="66667")
+        res = self.generic_serialization_test(
+            serializers.operation_serialization, no_test=True,
+            kwargs={"operation_queryset": result_queryset}
+        )
+        ope_json = json.loads(
+            res[('operations', 'archaeological_operations__Operation')]
+        )
+        self.assertEqual(len(ope_json), 1)
+        site_json = json.loads(
+            res[('operations', 'archaeological_operations__ArchaeologicalSite')]
+        )
+        self.assertEqual(len(site_json), 0)
 
     def test_conf_restore(self):
         current_number, zip_filename = self.generic_restore_test_genzip(