black shift
This commit is contained in:
@@ -63,7 +63,7 @@ class MongoDBHandler:
|
||||
if not hasattr(self, "_initialized") or not self._initialized:
|
||||
self._debug_mode = debug_mode
|
||||
self._mock_mode = mock_mode
|
||||
|
||||
|
||||
if mock_mode:
|
||||
# In mock mode, we don't need a real connection string
|
||||
self.uri = "mongodb://mock:27017/mockdb"
|
||||
@@ -76,7 +76,7 @@ class MongoDBHandler:
|
||||
# Use the configured connection string with authentication
|
||||
self.uri = mongo_configs.url
|
||||
print(f"Connecting to MongoDB: {self.uri}")
|
||||
|
||||
|
||||
# Define MongoDB client options with increased timeouts for better reliability
|
||||
self.client_options = {
|
||||
"maxPoolSize": 5,
|
||||
@@ -132,11 +132,13 @@ class CollectionContext:
|
||||
# If we're in mock mode, return a mock collection immediately
|
||||
if self.db_handler._mock_mode:
|
||||
return self._create_mock_collection()
|
||||
|
||||
|
||||
try:
|
||||
# Create a new client connection
|
||||
self.client = MongoClient(self.db_handler.uri, **self.db_handler.client_options)
|
||||
|
||||
self.client = MongoClient(
|
||||
self.db_handler.uri, **self.db_handler.client_options
|
||||
)
|
||||
|
||||
if self.db_handler._debug_mode:
|
||||
# In debug mode, we explicitly use the configured DB
|
||||
db_name = mongo_configs.DB
|
||||
@@ -148,7 +150,7 @@ class CollectionContext:
|
||||
except Exception:
|
||||
db_name = mongo_configs.DB
|
||||
print(f"Using fallback database '{db_name}'")
|
||||
|
||||
|
||||
self.collection = self.client[db_name][self.collection_name]
|
||||
|
||||
# Enhance collection methods with retry capabilities
|
||||
@@ -159,13 +161,17 @@ class CollectionContext:
|
||||
if "Authentication failed" in str(e):
|
||||
print(f"MongoDB authentication error: {e}")
|
||||
print("Attempting to reconnect with direct connection...")
|
||||
|
||||
|
||||
try:
|
||||
# Try a direct connection without authentication for testing
|
||||
direct_uri = f"mongodb://{mongo_configs.HOST}:{mongo_configs.PORT}/{mongo_configs.DB}"
|
||||
print(f"Trying direct connection: {direct_uri}")
|
||||
self.client = MongoClient(direct_uri, **self.db_handler.client_options)
|
||||
self.collection = self.client[mongo_configs.DB][self.collection_name]
|
||||
self.client = MongoClient(
|
||||
direct_uri, **self.db_handler.client_options
|
||||
)
|
||||
self.collection = self.client[mongo_configs.DB][
|
||||
self.collection_name
|
||||
]
|
||||
self._add_retry_capabilities()
|
||||
return self.collection
|
||||
except Exception as inner_e:
|
||||
@@ -181,54 +187,56 @@ class CollectionContext:
|
||||
if self.client:
|
||||
self.client.close()
|
||||
self.client = None
|
||||
|
||||
|
||||
return self._create_mock_collection()
|
||||
|
||||
|
||||
def _create_mock_collection(self):
|
||||
"""
|
||||
Create a mock collection for testing or graceful degradation.
|
||||
This prevents the application from crashing when MongoDB is unavailable.
|
||||
|
||||
|
||||
Returns:
|
||||
A mock MongoDB collection with simulated behaviors
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
||||
if self.db_handler._mock_mode:
|
||||
print(f"MOCK MODE: Using mock collection '{self.collection_name}'")
|
||||
else:
|
||||
print(f"Using mock MongoDB collection '{self.collection_name}' for graceful degradation")
|
||||
|
||||
print(
|
||||
f"Using mock MongoDB collection '{self.collection_name}' for graceful degradation"
|
||||
)
|
||||
|
||||
# Create in-memory storage for this mock collection
|
||||
if not hasattr(self.db_handler, '_mock_storage'):
|
||||
if not hasattr(self.db_handler, "_mock_storage"):
|
||||
self.db_handler._mock_storage = {}
|
||||
|
||||
|
||||
if self.collection_name not in self.db_handler._mock_storage:
|
||||
self.db_handler._mock_storage[self.collection_name] = []
|
||||
|
||||
|
||||
mock_collection = MagicMock()
|
||||
mock_data = self.db_handler._mock_storage[self.collection_name]
|
||||
|
||||
|
||||
# Define behavior for find operations
|
||||
def mock_find(query=None, *args, **kwargs):
|
||||
# Simple implementation that returns all documents
|
||||
return mock_data
|
||||
|
||||
|
||||
def mock_find_one(query=None, *args, **kwargs):
|
||||
# Simple implementation that returns the first matching document
|
||||
if not mock_data:
|
||||
return None
|
||||
return mock_data[0]
|
||||
|
||||
|
||||
def mock_insert_one(document, *args, **kwargs):
|
||||
# Add _id if not present
|
||||
if '_id' not in document:
|
||||
document['_id'] = f"mock_id_{len(mock_data)}"
|
||||
if "_id" not in document:
|
||||
document["_id"] = f"mock_id_{len(mock_data)}"
|
||||
mock_data.append(document)
|
||||
result = MagicMock()
|
||||
result.inserted_id = document['_id']
|
||||
result.inserted_id = document["_id"]
|
||||
return result
|
||||
|
||||
|
||||
def mock_insert_many(documents, *args, **kwargs):
|
||||
inserted_ids = []
|
||||
for doc in documents:
|
||||
@@ -237,40 +245,40 @@ class CollectionContext:
|
||||
result = MagicMock()
|
||||
result.inserted_ids = inserted_ids
|
||||
return result
|
||||
|
||||
|
||||
def mock_update_one(query, update, *args, **kwargs):
|
||||
result = MagicMock()
|
||||
result.modified_count = 1
|
||||
return result
|
||||
|
||||
|
||||
def mock_update_many(query, update, *args, **kwargs):
|
||||
result = MagicMock()
|
||||
result.modified_count = len(mock_data)
|
||||
return result
|
||||
|
||||
|
||||
def mock_delete_one(query, *args, **kwargs):
|
||||
result = MagicMock()
|
||||
result.deleted_count = 1
|
||||
if mock_data:
|
||||
mock_data.pop(0) # Just remove the first item for simplicity
|
||||
return result
|
||||
|
||||
|
||||
def mock_delete_many(query, *args, **kwargs):
|
||||
count = len(mock_data)
|
||||
mock_data.clear()
|
||||
result = MagicMock()
|
||||
result.deleted_count = count
|
||||
return result
|
||||
|
||||
|
||||
def mock_count_documents(query, *args, **kwargs):
|
||||
return len(mock_data)
|
||||
|
||||
|
||||
def mock_aggregate(pipeline, *args, **kwargs):
|
||||
return []
|
||||
|
||||
|
||||
def mock_create_index(keys, **kwargs):
|
||||
return f"mock_index_{keys}"
|
||||
|
||||
|
||||
# Assign the mock implementations
|
||||
mock_collection.find.side_effect = mock_find
|
||||
mock_collection.find_one.side_effect = mock_find_one
|
||||
@@ -283,10 +291,10 @@ class CollectionContext:
|
||||
mock_collection.count_documents.side_effect = mock_count_documents
|
||||
mock_collection.aggregate.side_effect = mock_aggregate
|
||||
mock_collection.create_index.side_effect = mock_create_index
|
||||
|
||||
|
||||
# Add retry capabilities to the mock collection
|
||||
self._add_retry_capabilities_to_mock(mock_collection)
|
||||
|
||||
|
||||
self.collection = mock_collection
|
||||
return self.collection
|
||||
|
||||
@@ -322,17 +330,25 @@ class CollectionContext:
|
||||
"""
|
||||
Add retry capabilities to mock collection methods.
|
||||
This is a simplified version that just wraps the mock methods.
|
||||
|
||||
|
||||
Args:
|
||||
mock_collection: The mock collection to enhance
|
||||
"""
|
||||
# List of common MongoDB collection methods to add retry capabilities to
|
||||
methods = [
|
||||
'insert_one', 'insert_many', 'find_one', 'find',
|
||||
'update_one', 'update_many', 'delete_one', 'delete_many',
|
||||
'replace_one', 'count_documents', 'aggregate'
|
||||
"insert_one",
|
||||
"insert_many",
|
||||
"find_one",
|
||||
"find",
|
||||
"update_one",
|
||||
"update_many",
|
||||
"delete_one",
|
||||
"delete_many",
|
||||
"replace_one",
|
||||
"count_documents",
|
||||
"aggregate",
|
||||
]
|
||||
|
||||
|
||||
# Add retry decorator to each method
|
||||
for method_name in methods:
|
||||
if hasattr(mock_collection, method_name):
|
||||
@@ -340,7 +356,7 @@ class CollectionContext:
|
||||
setattr(
|
||||
mock_collection,
|
||||
method_name,
|
||||
retry_operation(max_retries=1, retry_interval=0)(original_method)
|
||||
retry_operation(max_retries=1, retry_interval=0)(original_method),
|
||||
)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
|
||||
Reference in New Issue
Block a user