mirror of
https://github.com/langgenius/dify.git
synced 2025-08-15 12:47:08 +00:00
60 lines
2.3 KiB
Python
60 lines
2.3 KiB
Python
from unittest.mock import MagicMock
|
|
|
|
|
|
class ServiceDbTestHelper:
|
|
"""
|
|
Helper class for service database query tests.
|
|
"""
|
|
|
|
@staticmethod
|
|
def setup_db_query_filter_by_mock(mock_db, query_results):
|
|
"""
|
|
Smart database query mock that responds based on model type and query parameters.
|
|
|
|
Args:
|
|
mock_db: Mock database session
|
|
query_results: Dict mapping (model_name, filter_key, filter_value) to return value
|
|
Example: {('Account', 'email', 'test@example.com'): mock_account}
|
|
"""
|
|
|
|
def query_side_effect(model):
|
|
mock_query = MagicMock()
|
|
|
|
def filter_by_side_effect(**kwargs):
|
|
mock_filter_result = MagicMock()
|
|
|
|
def first_side_effect():
|
|
# Find matching result based on model and filter parameters
|
|
for (model_name, filter_key, filter_value), result in query_results.items():
|
|
if model.__name__ == model_name and filter_key in kwargs and kwargs[filter_key] == filter_value:
|
|
return result
|
|
return None
|
|
|
|
mock_filter_result.first.side_effect = first_side_effect
|
|
|
|
# Handle order_by calls for complex queries
|
|
def order_by_side_effect(*args, **kwargs):
|
|
mock_order_result = MagicMock()
|
|
|
|
def order_first_side_effect():
|
|
# Look for order_by results in the same query_results dict
|
|
for (model_name, filter_key, filter_value), result in query_results.items():
|
|
if (
|
|
model.__name__ == model_name
|
|
and filter_key == "order_by"
|
|
and filter_value == "first_available"
|
|
):
|
|
return result
|
|
return None
|
|
|
|
mock_order_result.first.side_effect = order_first_side_effect
|
|
return mock_order_result
|
|
|
|
mock_filter_result.order_by.side_effect = order_by_side_effect
|
|
return mock_filter_result
|
|
|
|
mock_query.filter_by.side_effect = filter_by_side_effect
|
|
return mock_query
|
|
|
|
mock_db.session.query.side_effect = query_side_effect
|