Add set method to MockSet

stphivos · Sep 8, 2024 · ec67f77 · ec67f77
1 parent 26257c5
commit ec67f77
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/django_mock_queries/query.py b/django_mock_queries/query.py
@@ -183,6 +183,10 @@ def distinct(self, *fields):
                 results[key] = item
         return self._mockset_class()(*results.values(), clone=self)
 
+    def set(self, objs, **attrs):
+        self.delete(**attrs)
+        self.add(*objs)
+
     def _raise_does_not_exist(self):
         does_not_exist = getattr(self.model, 'DoesNotExist', ObjectDoesNotExist)
         raise does_not_exist()

diff --git a/tests/test_query.py b/tests/test_query.py
@@ -1383,3 +1383,13 @@ def test_mock_set_annotation_by_nested_mock_model(self):
         )
         field1 = mockset.annotate(field1=models.F("nested_mock__field1")).values_list("field1")[0][0]
         assert field1 == "field_value"
+
+    def test_set_replaces_all_items(self):
+        mockset = MockSet(
+            MockModel(id=1, field="value_1", mock_name="item1"),
+            MockModel(id=2, field="value_2", mock_name="item2"),
+        )
+        mockset.set([MockModel(id=3, field="value_3", mock_name="item3")])
+
+        assert len(mockset) == 1
+        assert mockset[0].id == 3