MockSet.annotate function: Handle case when row._annotated_fields is …

…None (#175)
stphivos · Jun 27, 2024 · d38ad29 · d38ad29
2 parents 267aa1f + 5bfc9a1
commit d38ad29
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/django_mock_queries/query.py b/django_mock_queries/query.py
@@ -122,7 +122,7 @@ def annotate(self, **kwargs):
         results = list(self.items)
         for key, value in kwargs.items():
             for row in results:
-                if not hasattr(row, '_annotated_fields'):
+                if not (hasattr(row, '_annotated_fields') and isinstance(row._annotated_fields, list)):
                     row._annotated_fields = []
                 row._annotated_fields.append(key)
                 setattr(row, key, get_attribute(row, value)[0])

diff --git a/tests/test_query.py b/tests/test_query.py
@@ -1356,3 +1356,10 @@ def test_empty_queryset_filter(self):
         mockset = MockSet(car1, car2)
         self.assertEqual(mockset.count(), 2)
         self.assertEqual(mockset.filter(Q()).count(), 2)
+
+    def test_mock_set_annotation_by_nested_mock_model(self):
+        mockset = MockSet(
+            MockModel(id=1, nested_mock=MockModel(id=1, field1="field_value"))
+        )
+        field1 = mockset.annotate(field1=models.F("nested_mock__field1")).values_list("field1")[0][0]
+        assert field1 == "field_value"