diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py index 8d42d15b..e5ac2756 100644 --- a/tests/document/test_instance.py +++ b/tests/document/test_instance.py @@ -501,7 +501,7 @@ class TestDocumentInstance(MongoDBTestCase): doc.reload() Animal.drop_collection() - def test_update_shard_key_routing(self): + def test_save_update_shard_key_routing(self): """Ensures updating a doc with a specified shard_key includes it in the query. """ @@ -529,6 +529,36 @@ class TestDocumentInstance(MongoDBTestCase): Animal.drop_collection() + def test_save_create_shard_key_routing(self): + """Ensures inserting a doc with a specified shard_key includes it in + the query. + """ + + class Animal(Document): + _id = UUIDField(binary=False, primary_key=True, default=uuid.uuid4) + is_mammal = BooleanField() + name = StringField() + meta = {"shard_key": ("is_mammal",)} + + Animal.drop_collection() + doc = Animal(is_mammal=True, name="Dog") + + mongo_db = get_mongodb_version() + + with query_counter() as q: + doc.save() + query_op = q.db.system.profile.find({"ns": "mongoenginetest.animal"})[0] + assert query_op["op"] == "command" + assert query_op["command"]["findAndModify"] == "animal" + if mongo_db <= MONGODB_34: + assert set(query_op["query"].keys()) == set(["_id", "is_mammal"]) + else: + assert set(query_op["command"]["query"].keys()) == set( + ["_id", "is_mammal"] + ) + + Animal.drop_collection() + def test_reload_with_changed_fields(self): """Ensures reloading will not affect changed fields"""