From 78b240b740b34de450d30a00c669a9283a8b37de Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastien=20G=C3=A9rard?= <bastien.gerard@cluepoints.com>
Date: Wed, 4 Dec 2019 21:49:17 +0100
Subject: [PATCH] updated changelog + improved query_counter test

---
 docs/changelog.rst              |  1 +
 mongoengine/context_managers.py |  2 +-
 tests/document/test_instance.py | 38 ---------------------------
 tests/test_context_managers.py  | 46 +++++++++++++++++++++++++++++++++
 4 files changed, 48 insertions(+), 39 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index 102e826d..99081957 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -24,6 +24,7 @@ Development
 - The codebase is now formatted using ``black``. #2109
 - In bulk write insert, the detailed error message would raise in exception.
 - Added ability to compare Q and Q operations #2204
+- Added ability to use a db alias on query_counter #2194
 
 Changes in 0.18.2
 =================
diff --git a/mongoengine/context_managers.py b/mongoengine/context_managers.py
index 5920b724..1592ceef 100644
--- a/mongoengine/context_managers.py
+++ b/mongoengine/context_managers.py
@@ -171,7 +171,7 @@ class no_sub_classes(object):
 class query_counter(object):
     """Query_counter context manager to get the number of queries.
     This works by updating the `profiling_level` of the database so that all queries get logged,
-    resetting the db.system.profile collection at the beginnig of the context and counting the new entries.
+    resetting the db.system.profile collection at the beginning of the context and counting the new entries.
 
     This was designed for debugging purpose. In fact it is a global counter so queries issued by other threads/processes
     can interfere with it
diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py
index c8ad2ff3..173e02f2 100644
--- a/tests/document/test_instance.py
+++ b/tests/document/test_instance.py
@@ -2825,44 +2825,6 @@ class TestInstance(MongoDBTestCase):
 
         assert "testdb-1" == B._meta.get("db_alias")
 
-    def test_query_counter_alias(self):
-        """query_counter works properly with db aliases?"""
-        # Register a connection with db_alias testdb-1
-        register_connection("testdb-1", "mongoenginetest2")
-
-        class A(Document):
-            """Uses default db_alias
-            """
-
-            name = StringField()
-
-        class B(Document):
-            """Uses testdb-1 db_alias
-            """
-
-            name = StringField()
-            meta = {"db_alias": "testdb-1"}
-
-        with query_counter() as q:
-            assert q == 0
-            a = A.objects.create(name="A")
-            assert q == 1
-            a = A.objects.first()
-            assert q == 2
-            a.name = "Test A"
-            a.save()
-            assert q == 3
-
-        with query_counter(alias="testdb-1") as q:
-            assert q == 0
-            b = B.objects.create(name="B")
-            assert q == 1
-            b = B.objects.first()
-            assert q == 2
-            b.name = "Test B"
-            b.save()
-            assert q == 3
-
     def test_db_ref_usage(self):
         """DB Ref usage in dict_fields."""
 
diff --git a/tests/test_context_managers.py b/tests/test_context_managers.py
index c10a0224..fa3f5960 100644
--- a/tests/test_context_managers.py
+++ b/tests/test_context_managers.py
@@ -282,6 +282,52 @@ class TestContextManagers:
             assert q < 1000
             assert q <= int(q)
 
+    def test_query_counter_alias(self):
+        """query_counter works properly with db aliases?"""
+        # Register a connection with db_alias testdb-1
+        register_connection("testdb-1", "mongoenginetest2")
+
+        class A(Document):
+            """Uses default db_alias"""
+
+            name = StringField()
+
+        class B(Document):
+            """Uses testdb-1 db_alias"""
+
+            name = StringField()
+            meta = {"db_alias": "testdb-1"}
+
+        A.drop_collection()
+        B.drop_collection()
+
+        with query_counter() as q:
+            assert q == 0
+            A.objects.create(name="A")
+            assert q == 1
+            a = A.objects.first()
+            assert q == 2
+            a.name = "Test A"
+            a.save()
+            assert q == 3
+            # querying the other db should'nt alter the counter
+            B.objects().first()
+            assert q == 3
+
+        with query_counter(alias="testdb-1") as q:
+            assert q == 0
+            B.objects.create(name="B")
+            assert q == 1
+            b = B.objects.first()
+            assert q == 2
+            b.name = "Test B"
+            b.save()
+            assert b.name == "Test B"
+            assert q == 3
+            # querying the other db should'nt alter the counter
+            A.objects().first()
+            assert q == 3
+
     def test_query_counter_counts_getmore_queries(self):
         connect("mongoenginetest")
         db = get_db()