Add count return value to anonymise_queryset

yunojuno · Sep 10, 2023 · 7cf6869 · 7cf6869
1 parent 0f357a9
commit 7cf6869
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/anonymiser/models.py b/anonymiser/models.py
@@ -115,11 +115,14 @@ def anonymise_object(self, obj: models.Model) -> None:
             output[field.name] = self.anonymise_field(obj, field.name)
         self.post_anonymise_object(obj, **output)
 
-    def anonymise_queryset(self, queryset: models.QuerySet) -> None:
+    def anonymise_queryset(self, queryset: models.QuerySet) -> int:
         """Anonymise all objects in the queryset (and SAVE)."""
+        count = 0
         for obj in queryset:
             self.anonymise_object(obj)
             obj.save()
+            count += 1
+        return count
 
     def post_anonymise_object(
         self, obj: models.Model, **updates: AnonymisationResult

diff --git a/tests/test_models.py b/tests/test_models.py
@@ -90,3 +90,9 @@ def test_get_anonymisable_fields(self, user_anonymiser: UserAnonymiser) -> None:
         assert user_anonymiser.get_anonymisable_fields() == [
             User._meta.get_field("first_name")
         ]
+
+    def test_anonymise_queryset(
+        self, user: User, user_anonymiser: UserAnonymiser
+    ) -> None:
+        assert user_anonymiser.anonymise_queryset(User.objects.none()) == 0
+        assert user_anonymiser.anonymise_queryset(User.objects.all()) == 1