fix: avoid patching collections multiple times; closes #49

Author: Michael Kryukov

mongomock_motor/__init__.py

@@ -193,7 +193,7 @@ async def to_list(self, *args, **kwargs):
 class AsyncMongoMockCollection:
     def __init__(self, database, collection):
         self.database = database
-        self.__collection = collection
+        self.__collection = _patch_collection_internals(collection)
 
     def get_io_loop(self):
         return self.database.get_io_loop()
@@ -248,9 +248,7 @@ def get_io_loop(self):
     def get_collection(self, *args, **kwargs):
         return AsyncMongoMockCollection(
             self,
-            _patch_collection_internals(
-                self.__database.get_collection(*args, **kwargs),
-            ),
+            self.__database.get_collection(*args, **kwargs),
         )
 
     def aggregate(self, *args, **kwargs) -> AsyncCursor:

mongomock_motor/patches.py

@@ -61,9 +61,11 @@ def wrapper(data, *args, **kwargs):
 
         return wrapper
 
-    collection._insert = with_enriched_duplicate_key_error(collection._insert)
+    collection._insert = with_enriched_duplicate_key_error(
+        collection._insert,
+    )
     collection._ensure_uniques = with_enriched_duplicate_key_error(
-        collection._ensure_uniques
+        collection._ensure_uniques,
     )
 
     return collection
@@ -89,19 +91,27 @@ def _patch_iter_documents(collection):
     that is inherited from "str". Looks like pymongo works ok
     with that, so we should be too.
     """
-    _iter_documents = collection._iter_documents
 
-    def iter_documents(filter):
-        return _iter_documents(_normalize_strings(filter))
+    def with_normalized_strings_in_filter(fn):
+        @wraps(fn)
+        def wrapper(filter):
+            return fn(_normalize_strings(filter))
 
-    collection._iter_documents = iter_documents
+        return wrapper
+
+    collection._iter_documents = with_normalized_strings_in_filter(
+        collection._iter_documents,
+    )
 
     return collection
 
 
 def _patch_collection_internals(collection):
+    if getattr(collection, '_patched_by_mongomock_motor', False):
+        return collection
     collection = _patch_insert_and_ensure_uniques(collection)
     collection = _patch_iter_documents(collection)
+    collection._patched_by_mongomock_motor = True
     return collection
 
 

tests/test_mocking.py

@@ -2,9 +2,11 @@
 
 import pytest
 from bson import ObjectId
+from pymongo.read_preferences import Primary
 from pymongo.results import UpdateResult
 
 from mongomock_motor import AsyncMongoMockClient
+from mongomock_motor.patches import _patch_iter_documents
 
 
 @pytest.mark.anyio
@@ -47,3 +49,33 @@ async def test_patch_object():
     ):
         with pytest.raises(RuntimeError):
             await sample_function(collection)
+
+
+@pytest.mark.anyio
+async def test_no_multiple_patching():
+    database = AsyncMongoMockClient()['test']
+
+    with patch(
+        'mongomock_motor.patches._patch_iter_documents',
+        wraps=_patch_iter_documents,
+    ) as patch_iter_documents:
+        for _ in range(2):
+            collection = database['test']
+            assert collection
+
+        assert patch_iter_documents.call_count == 1
+
+        for _ in range(2):
+            collection = database.get_collection(
+                'test',
+                read_preference=Primary,
+            )
+            assert collection
+
+        assert patch_iter_documents.call_count == 3
+
+        for _ in range(2):
+            collection = database.get_collection('test')
+            assert collection
+
+        assert patch_iter_documents.call_count == 3