Merge pull request #169 from harry-kim/fix/in_bulk

Fix in_bulk to use field_name
stphivos · Jun 1, 2023 · 65ef8b8 · 65ef8b8
2 parents c6c838c + ea34362
commit 65ef8b8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/django_mock_queries/query.py b/django_mock_queries/query.py
@@ -114,7 +114,7 @@ def exists(self):
     def in_bulk(self, id_list=None, *, field_name='pk'):
         result = {}
         for model in self.items:
-            if id_list is None or model.pk in id_list:
+            if id_list is None or getattr(model, field_name) in id_list:
                 result[getattr(model, field_name)] = model
         return result
 

diff --git a/tests/test_query.py b/tests/test_query.py
@@ -1024,7 +1024,7 @@ def test_in_bulk(self):
         qs = MockSet(golf, polo, kia)
 
         self.assertEqual(qs.in_bulk(), {1: golf, 2: polo, 4: kia})
-        self.assertEqual(qs.in_bulk(id_list=[4], field_name='model'), {'kia': kia})
+        self.assertEqual(qs.in_bulk(id_list=['kia'], field_name='model'), {'kia': kia})
 
     def test_annotate(self):
         qs = MockSet(CarVariation(color='green', car=Car(model='golf', id=1), id=1),