From b407c0e6c6a812b80ba3882bea0d9dd1ecac4ad4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Mon, 4 Mar 2019 23:42:45 +0100
Subject: [PATCH] add test for shard key routing (ported from
 https://github.com/closeio/mongoengine/commit/43f35f5)

---
 tests/document/instance.py | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/tests/document/instance.py b/tests/document/instance.py
index f4527ad8..04bddea1 100644
--- a/tests/document/instance.py
+++ b/tests/document/instance.py
@@ -466,7 +466,13 @@ class InstanceTest(MongoDBTestCase):
         Animal.drop_collection()
         doc = Animal(superphylum='Deuterostomia')
         doc.save()
-        doc.reload()
+
+        with query_counter() as q:
+            doc.reload()
+            query_op = q.db.system.profile.find({'ns': 'mongoenginetest.animal'})[0]
+            self.assertEqual(set(query_op['query']['filter'].keys()), set(['_id', 'superphylum']))
+
+        Animal.drop_collection()
 
     def test_reload_sharded_nested(self):
         class SuperPhylum(EmbeddedDocument):
@@ -480,6 +486,29 @@ class InstanceTest(MongoDBTestCase):
         doc = Animal(superphylum=SuperPhylum(name='Deuterostomia'))
         doc.save()
         doc.reload()
+        Animal.drop_collection()
+
+    def test_update_shard_key_routing(self):
+        """Ensures updating a doc with a specified shard_key includes it in
+        the query.
+        """
+        class Animal(Document):
+            is_mammal = BooleanField()
+            name = StringField()
+            meta = {'shard_key': ('is_mammal', 'id')}
+
+        Animal.drop_collection()
+        doc = Animal(is_mammal=True, name='Dog')
+        doc.save()
+
+        with query_counter() as q:
+            doc.name = 'Cat'
+            doc.save()
+            query_op = q.db.system.profile.find({'ns': 'mongoenginetest.animal'})[0]
+            self.assertEqual(query_op['op'], 'update')
+            self.assertEqual(set(query_op['query'].keys()), set(['_id', 'is_mammal']))
+
+        Animal.drop_collection()
 
     def test_reload_with_changed_fields(self):
         """Ensures reloading will not affect changed fields"""