dify/api/tests/unit_tests/services/services_test_help.py

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