diff --git a/django_mock_queries/query.py b/django_mock_queries/query.py index 93d2ced..2cf78c9 100644 --- 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 index c11f4e5..66ad52a 100644 --- 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