cleaning test pymongo clients properly

Author: shaypal5

tests/conftest.py

@@ -107,14 +107,32 @@ def cleanup_mongo_clients():
     yield
 
     # Cleanup after all tests
+    import contextlib
+
     try:
-        from tests.test_mongo_core import _test_mongetter
+        from tests.test_mongo_core import _mongo_clients, _test_mongetter
+
+        # Close all tracked MongoDB clients
+        for client in _mongo_clients:
+            with contextlib.suppress(Exception):
+                client.close()
+
+        # Clear the list for next test run
+        _mongo_clients.clear()
 
+        # Also clean up _test_mongetter specifically
         if hasattr(_test_mongetter, "client"):
-            # Close the MongoDB client to avoid ResourceWarning
-            _test_mongetter.client.close()
             # Remove the client attribute so future test runs start fresh
             delattr(_test_mongetter, "client")
+
+        # Clean up any _custom_mongetter functions that may have been created
+        import tests.test_mongo_core
+
+        for attr_name in dir(tests.test_mongo_core):
+            attr = getattr(tests.test_mongo_core, attr_name)
+            if callable(attr) and hasattr(attr, "client"):
+                delattr(attr, "client")
+
     except (ImportError, AttributeError):
         # If the module wasn't imported or client wasn't created,
         # then there's nothing to clean up

tests/test_mongo_core.py

@@ -107,14 +107,46 @@ def _get_cachier_db_mongo_client():
 )
 
 
+# Global registry to track all MongoDB clients created during tests
+_mongo_clients = []
+
+
+def cleanup_all_mongo_clients():
+    """Clean up all MongoDB clients to prevent ResourceWarning."""
+    import contextlib
+    import sys
+
+    global _mongo_clients
+
+    # Close all tracked clients
+    for client in _mongo_clients:
+        with contextlib.suppress(Exception):
+            client.close()
+
+    # Clear the list
+    _mongo_clients.clear()
+
+    # Clean up any mongetter functions with clients
+    current_module = sys.modules[__name__]
+    for attr_name in dir(current_module):
+        attr = getattr(current_module, attr_name)
+        if callable(attr) and hasattr(attr, "client"):
+            with contextlib.suppress(Exception):
+                if hasattr(attr.client, "close"):
+                    attr.client.close()
+                delattr(attr, "client")
+
+
 def _test_mongetter():
     if not hasattr(_test_mongetter, "client"):
         if str(CFG.mget(CfgKey.TEST_VS_DOCKERIZED_MONGO)).lower() == "true":
             print("Using live MongoDB instance for testing.")
             _test_mongetter.client = _get_cachier_db_mongo_client()
+            _mongo_clients.append(_test_mongetter.client)
         else:
             print("Using in-memory MongoDB instance for testing.")
             _test_mongetter.client = InMemoryMongoClient()
+            _mongo_clients.append(_test_mongetter.client)
     db_obj = _test_mongetter.client["cachier_test"]
     if _COLLECTION_NAME not in db_obj.list_collection_names():
         db_obj.create_collection(_COLLECTION_NAME)
@@ -137,17 +169,29 @@ def _custom_mongetter():
             ):
                 print("Using live MongoDB instance for testing.")
                 _custom_mongetter.client = _get_cachier_db_mongo_client()
+                _mongo_clients.append(_custom_mongetter.client)
             else:
                 print("Using in-memory MongoDB instance for testing.")
                 _custom_mongetter.client = InMemoryMongoClient()
+                _mongo_clients.append(_custom_mongetter.client)
         db_obj = _custom_mongetter.client["cachier_test"]
         if _COLLECTION_NAME not in db_obj.list_collection_names():
             db_obj.create_collection(collection_name)
         return db_obj[collection_name]
 
+    # Store the mongetter function for cleanup
+    _custom_mongetter._collection_name = collection_name
     return _custom_mongetter
 
 
+@pytest.fixture(autouse=True)
+def mongo_cleanup():
+    """Ensure MongoDB clients are cleaned up after each test."""
+    yield
+    # Clean up after test
+    cleanup_all_mongo_clients()
+
+
 # === Mongo core tests ===