refactoring in tests

Author: bagerard

tests/fields/test_generic_reference_field.py

@@ -10,12 +10,12 @@
     ValidationError,
 )
 from mongoengine.base import _DocumentRegistry
-from tests.utils import MongoDBTestCase
+from tests.utils import MongoDBTestCase, get_as_pymongo
 
 
 class TestField(MongoDBTestCase):
 
-    def test_generic_reference(self):
+    def test_generic_reference_field_basics(self):
         """Ensure that a GenericReferenceField properly dereferences items."""
 
         class Link(Document):
@@ -42,14 +42,27 @@ class Bookmark(Document):
         bm.save()
 
         bm = Bookmark.objects(bookmark_object=post_1).first()
-
+        assert get_as_pymongo(bm) == {
+            "_id": bm.id,
+            "bookmark_object": {
+                "_cls": "Post",
+                "_ref": post_1.to_dbref(),
+            },
+        }
         assert bm.bookmark_object == post_1
         assert isinstance(bm.bookmark_object, Post)
 
         bm.bookmark_object = link_1
         bm.save()
 
         bm = Bookmark.objects(bookmark_object=link_1).first()
+        assert get_as_pymongo(bm) == {
+            "_id": bm.id,
+            "bookmark_object": {
+                "_cls": "Link",
+                "_ref": link_1.to_dbref(),
+            },
+        }
 
         assert bm.bookmark_object == link_1
         assert isinstance(bm.bookmark_object, Link)
@@ -67,14 +80,6 @@ class OtherObj(Document):
         s1 = SomeObj().save()
         OtherObj(obj=s1).save()
 
-        assert OtherObj.objects.as_pymongo().first() == {
-            "_id": OtherObj.objects.first().id,
-            "obj": {
-                "_cls": "SomeObj",
-                "_ref": s1.to_dbref(),
-            },
-        }
-
         # Query using to_dbref
         assert OtherObj.objects(obj__in=[s1.to_dbref()]).count() == 1