diff options
Diffstat (limited to 'ishtar_common')
| -rw-r--r-- | ishtar_common/serializers.py | 15 | ||||
| -rw-r--r-- | ishtar_common/tests.py | 17 |
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): |