Serializer: manage queryset filter

2 files changed, 23 insertions, 9 deletions

diff --git a/ishtar_common/serializers.py b/ishtar_common/serializers.py
index 03433df89..409051ec0 100644
--- a/ishtar_common/serializers.py
+++ b/ishtar_common/serializers.py
@@ -103,13 +103,18 @@ def archive_serialization(result, archive_dir=None, archive=False,
     return archive_name
 
 
-def generic_get_results(model_list, dirname, no_geo=True):
+def generic_get_results(model_list, dirname, no_geo=True,
+                        result_queryset=None):
     result = OrderedDict()
     for model in model_list:
-        model_name = model.__name__
-        model_name = str(model.__module__).split(".")[0] + "__" + model_name
-
-        base_q = model.objects
+        base_model_name = model.__name__
+        model_name = str(model.__module__).split(".")[0] + "__" + \
+                     base_model_name
+
+        if result_queryset and base_model_name in result_queryset:
+            base_q = result_queryset[base_model_name]
+        else:
+            base_q = model.objects
         q = base_q
         recursion = None
         if hasattr(model, "parent"):
diff --git a/ishtar_common/tests.py b/ishtar_common/tests.py
index 7bb3ce098..68c206865 100644
--- a/ishtar_common/tests.py
+++ b/ishtar_common/tests.py
@@ -599,8 +599,12 @@ class CacheTest(TestCase):
 
 
 class GenericSerializationTest:
-    def generic_serialization_test(self, serialize):
-        json_result = serialize()
+    def generic_serialization_test(self, serialize, no_test=False, kwargs=None):
+        if not kwargs:
+            kwargs = {}
+        json_result = serialize(**kwargs)
+        if no_test:
+            return json_result
         for key in json_result.keys():
             __, k = key
             module_name, model_name = k.split("__")
@@ -615,12 +619,17 @@ class GenericSerializationTest:
                 msg="Serialization for model {}.{} failed. {} serialized {} "
                     "expected".format(module.__name__, model_name,
                                       serialization_count, current_count))
+        return json_result
 
-    def generic_restore_test_genzip(self, model_list, serialization):
+    def generic_restore_test_genzip(self, model_list, serialization,
+                                    kwargs=None):
         current_number = {}
         for model in model_list:
             current_number[model.__name__] = model.objects.count()
-        zip_filename = serialization(archive=True)
+        if not kwargs:
+            kwargs = {}
+        kwargs["archive"] = True
+        zip_filename = serialization(**kwargs)
         return current_number, zip_filename
 
     def generic_restore_test(self, zip_filename, current_number, model_list):