test(api): add autospec to MagicMock-based patch usage (#32752)

This commit is contained in:
-LAN-
2026-03-01 04:30:45 +08:00
committed by GitHub
parent c034eb036c
commit 20fcc95db9
86 changed files with 865 additions and 804 deletions

View File

@ -19,14 +19,14 @@ class TestAgentService:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.agent_service.PluginAgentClient") as mock_plugin_agent_client,
patch("services.agent_service.ToolManager") as mock_tool_manager,
patch("services.agent_service.AgentConfigManager") as mock_agent_config_manager,
patch("services.agent_service.PluginAgentClient", autospec=True) as mock_plugin_agent_client,
patch("services.agent_service.ToolManager", autospec=True) as mock_tool_manager,
patch("services.agent_service.AgentConfigManager", autospec=True) as mock_agent_config_manager,
patch("services.agent_service.current_user", create_autospec(Account, instance=True)) as mock_current_user,
patch("services.app_service.FeatureService") as mock_feature_service,
patch("services.app_service.EnterpriseService") as mock_enterprise_service,
patch("services.app_service.ModelManager") as mock_model_manager,
patch("services.account_service.FeatureService") as mock_account_feature_service,
patch("services.app_service.FeatureService", autospec=True) as mock_feature_service,
patch("services.app_service.EnterpriseService", autospec=True) as mock_enterprise_service,
patch("services.app_service.ModelManager", autospec=True) as mock_model_manager,
patch("services.account_service.FeatureService", autospec=True) as mock_account_feature_service,
):
# Setup default mock returns for agent service
mock_plugin_agent_client_instance = mock_plugin_agent_client.return_value

View File

@ -18,18 +18,22 @@ class TestAppGenerateService:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.billing_service.BillingService") as mock_billing_service,
patch("services.app_generate_service.WorkflowService") as mock_workflow_service,
patch("services.app_generate_service.RateLimit") as mock_rate_limit,
patch("services.app_generate_service.CompletionAppGenerator") as mock_completion_generator,
patch("services.app_generate_service.ChatAppGenerator") as mock_chat_generator,
patch("services.app_generate_service.AgentChatAppGenerator") as mock_agent_chat_generator,
patch("services.app_generate_service.AdvancedChatAppGenerator") as mock_advanced_chat_generator,
patch("services.app_generate_service.WorkflowAppGenerator") as mock_workflow_generator,
patch("services.app_generate_service.MessageBasedAppGenerator") as mock_message_based_generator,
patch("services.account_service.FeatureService") as mock_account_feature_service,
patch("services.app_generate_service.dify_config") as mock_dify_config,
patch("configs.dify_config") as mock_global_dify_config,
patch("services.billing_service.BillingService", autospec=True) as mock_billing_service,
patch("services.app_generate_service.WorkflowService", autospec=True) as mock_workflow_service,
patch("services.app_generate_service.RateLimit", autospec=True) as mock_rate_limit,
patch("services.app_generate_service.CompletionAppGenerator", autospec=True) as mock_completion_generator,
patch("services.app_generate_service.ChatAppGenerator", autospec=True) as mock_chat_generator,
patch("services.app_generate_service.AgentChatAppGenerator", autospec=True) as mock_agent_chat_generator,
patch(
"services.app_generate_service.AdvancedChatAppGenerator", autospec=True
) as mock_advanced_chat_generator,
patch("services.app_generate_service.WorkflowAppGenerator", autospec=True) as mock_workflow_generator,
patch(
"services.app_generate_service.MessageBasedAppGenerator", autospec=True
) as mock_message_based_generator,
patch("services.account_service.FeatureService", autospec=True) as mock_account_feature_service,
patch("services.app_generate_service.dify_config", autospec=True) as mock_dify_config,
patch("configs.dify_config", autospec=True) as mock_global_dify_config,
):
# Setup default mock returns for billing service
mock_billing_service.update_tenant_feature_plan_usage.return_value = {
@ -983,7 +987,7 @@ class TestAppGenerateService:
}
# Execute the method under test
with patch("services.app_generate_service.AppExecutionParams") as mock_exec_params:
with patch("services.app_generate_service.AppExecutionParams", autospec=True) as mock_exec_params:
mock_payload = MagicMock()
mock_payload.workflow_run_id = fake.uuid4()
mock_payload.model_dump_json.return_value = "{}"

View File

@ -17,10 +17,12 @@ class TestModelLoadBalancingService:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.model_load_balancing_service.ProviderManager") as mock_provider_manager,
patch("services.model_load_balancing_service.LBModelManager") as mock_lb_model_manager,
patch("services.model_load_balancing_service.ModelProviderFactory") as mock_model_provider_factory,
patch("services.model_load_balancing_service.encrypter") as mock_encrypter,
patch("services.model_load_balancing_service.ProviderManager", autospec=True) as mock_provider_manager,
patch("services.model_load_balancing_service.LBModelManager", autospec=True) as mock_lb_model_manager,
patch(
"services.model_load_balancing_service.ModelProviderFactory", autospec=True
) as mock_model_provider_factory,
patch("services.model_load_balancing_service.encrypter", autospec=True) as mock_encrypter,
):
# Setup default mock returns
mock_provider_manager_instance = mock_provider_manager.return_value

View File

@ -17,8 +17,8 @@ class TestModelProviderService:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.model_provider_service.ProviderManager") as mock_provider_manager,
patch("services.model_provider_service.ModelProviderFactory") as mock_model_provider_factory,
patch("services.model_provider_service.ProviderManager", autospec=True) as mock_provider_manager,
patch("services.model_provider_service.ModelProviderFactory", autospec=True) as mock_model_provider_factory,
):
# Setup default mock returns
mock_provider_manager.return_value.get_configurations.return_value = MagicMock()
@ -526,7 +526,9 @@ class TestModelProviderService:
# Act: Execute the method under test
service = ModelProviderService()
with patch.object(service, "get_provider_credential", return_value=expected_credentials) as mock_method:
with patch.object(
service, "get_provider_credential", return_value=expected_credentials, autospec=True
) as mock_method:
result = service.get_provider_credential(tenant.id, "openai")
# Assert: Verify the expected outcomes
@ -854,7 +856,9 @@ class TestModelProviderService:
# Act: Execute the method under test
service = ModelProviderService()
with patch.object(service, "get_model_credential", return_value=expected_credentials) as mock_method:
with patch.object(
service, "get_model_credential", return_value=expected_credentials, autospec=True
) as mock_method:
result = service.get_model_credential(tenant.id, "openai", "llm", "gpt-4", None)
# Assert: Verify the expected outcomes

View File

@ -22,16 +22,13 @@ class TestWebhookService:
def mock_external_dependencies(self):
"""Mock external service dependencies."""
with (
patch("services.trigger.webhook_service.AsyncWorkflowService") as mock_async_service,
patch("services.trigger.webhook_service.ToolFileManager") as mock_tool_file_manager,
patch("services.trigger.webhook_service.file_factory") as mock_file_factory,
patch("services.account_service.FeatureService") as mock_feature_service,
patch("services.trigger.webhook_service.AsyncWorkflowService", autospec=True) as mock_async_service,
patch("services.trigger.webhook_service.ToolFileManager", autospec=True) as mock_tool_file_manager,
patch("services.trigger.webhook_service.file_factory", autospec=True) as mock_file_factory,
patch("services.account_service.FeatureService", autospec=True) as mock_feature_service,
):
# Mock ToolFileManager
mock_tool_file_instance = MagicMock()
mock_tool_file_manager.return_value = mock_tool_file_instance
# Mock file creation
mock_tool_file_instance = mock_tool_file_manager.return_value # Mock file creation
mock_tool_file = MagicMock()
mock_tool_file.id = "test_file_id"
mock_tool_file_instance.create_file_by_raw.return_value = mock_tool_file
@ -435,12 +432,12 @@ class TestWebhookService:
with flask_app_with_containers.app_context():
# Mock tenant owner lookup to return the test account
with patch("services.trigger.webhook_service.select") as mock_select:
with patch("services.trigger.webhook_service.select", autospec=True) as mock_select:
mock_query = MagicMock()
mock_select.return_value.join.return_value.where.return_value = mock_query
# Mock the session to return our test account
with patch("services.trigger.webhook_service.Session") as mock_session:
with patch("services.trigger.webhook_service.Session", autospec=True) as mock_session:
mock_session_instance = MagicMock()
mock_session.return_value.__enter__.return_value = mock_session_instance
mock_session_instance.scalar.return_value = test_data["account"]
@ -462,7 +459,7 @@ class TestWebhookService:
with flask_app_with_containers.app_context():
# Mock EndUserService to raise an exception
with patch(
"services.trigger.webhook_service.EndUserService.get_or_create_end_user_by_type"
"services.trigger.webhook_service.EndUserService.get_or_create_end_user_by_type", autospec=True
) as mock_end_user:
mock_end_user.side_effect = ValueError("Failed to create end user")

View File

@ -764,7 +764,7 @@ class TestWorkflowService:
# Act - Mock current_user context and pass session
from unittest.mock import patch
with patch("flask_login.utils._get_user", return_value=account):
with patch("flask_login.utils._get_user", return_value=account, autospec=True):
result = workflow_service.publish_workflow(
session=db_session_with_containers, app_model=app, account=account
)
@ -1401,6 +1401,7 @@ class TestWorkflowService:
DifyNodeFactory,
"_build_model_instance_for_llm_node",
return_value=MagicMock(spec=ModelInstance),
autospec=True,
):
result = workflow_service.run_free_workflow_node(
node_data=node_data, tenant_id=tenant_id, user_id=user_id, node_id=node_id, user_inputs=user_inputs

View File

@ -18,7 +18,9 @@ class TestAddDocumentToIndexTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.add_document_to_index_task.IndexProcessorFactory") as mock_index_processor_factory,
patch(
"tasks.add_document_to_index_task.IndexProcessorFactory", autospec=True
) as mock_index_processor_factory,
):
# Setup mock index processor
mock_processor = MagicMock()
@ -378,7 +380,7 @@ class TestAddDocumentToIndexTask:
redis_client.set(indexing_cache_key, "processing", ex=300)
# Mock the get_child_chunks method for each segment
with patch.object(DocumentSegment, "get_child_chunks") as mock_get_child_chunks:
with patch.object(DocumentSegment, "get_child_chunks", autospec=True) as mock_get_child_chunks:
# Setup mock to return child chunks for each segment
mock_child_chunks = []
for i in range(2): # Each segment has 2 child chunks

View File

@ -51,9 +51,9 @@ class TestBatchCreateSegmentToIndexTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.batch_create_segment_to_index_task.storage") as mock_storage,
patch("tasks.batch_create_segment_to_index_task.ModelManager") as mock_model_manager,
patch("tasks.batch_create_segment_to_index_task.VectorService") as mock_vector_service,
patch("tasks.batch_create_segment_to_index_task.storage", autospec=True) as mock_storage,
patch("tasks.batch_create_segment_to_index_task.ModelManager", autospec=True) as mock_model_manager,
patch("tasks.batch_create_segment_to_index_task.VectorService", autospec=True) as mock_vector_service,
):
# Setup default mock returns
mock_storage.download.return_value = None

View File

@ -63,8 +63,8 @@ class TestCleanDatasetTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.clean_dataset_task.storage") as mock_storage,
patch("tasks.clean_dataset_task.IndexProcessorFactory") as mock_index_processor_factory,
patch("tasks.clean_dataset_task.storage", autospec=True) as mock_storage,
patch("tasks.clean_dataset_task.IndexProcessorFactory", autospec=True) as mock_index_processor_factory,
):
# Setup default mock returns
mock_storage.delete.return_value = None
@ -597,7 +597,7 @@ class TestCleanDatasetTask:
db_session_with_containers.commit()
# Mock the get_image_upload_file_ids function to return our image file IDs
with patch("tasks.clean_dataset_task.get_image_upload_file_ids") as mock_get_image_ids:
with patch("tasks.clean_dataset_task.get_image_upload_file_ids", autospec=True) as mock_get_image_ids:
mock_get_image_ids.return_value = [f.id for f in image_files]
# Execute the task

View File

@ -41,7 +41,7 @@ class TestCreateSegmentToIndexTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.create_segment_to_index_task.IndexProcessorFactory") as mock_factory,
patch("tasks.create_segment_to_index_task.IndexProcessorFactory", autospec=True) as mock_factory,
):
# Setup default mock returns
mock_processor = MagicMock()
@ -708,7 +708,7 @@ class TestCreateSegmentToIndexTask:
redis_client.set(cache_key, "processing", ex=300)
# Mock Redis to raise exception in finally block
with patch.object(redis_client, "delete", side_effect=Exception("Redis connection failed")):
with patch.object(redis_client, "delete", side_effect=Exception("Redis connection failed"), autospec=True):
# Act: Execute the task - Redis failure should not prevent completion
with pytest.raises(Exception) as exc_info:
create_segment_to_index_task(segment.id)

View File

@ -37,7 +37,7 @@ class _TrackedSessionContext:
self._closed_sessions.append(self._session)
return original_close(*args, **kwargs)
self._close_patcher = patch.object(self._session, "close", side_effect=_tracked_close)
self._close_patcher = patch.object(self._session, "close", side_effect=_tracked_close, autospec=True)
self._close_patcher.start()
return self._session
@ -69,7 +69,9 @@ def session_close_tracker():
original_context_manager = original_create_session(*args, **kwargs)
return _TrackedSessionContext(original_context_manager, opened_sessions, closed_sessions)
with patch.object(task_module.session_factory, "create_session", side_effect=_tracked_create_session):
with patch.object(
task_module.session_factory, "create_session", side_effect=_tracked_create_session, autospec=True
):
yield {"opened_sessions": opened_sessions, "closed_sessions": closed_sessions}
@ -77,13 +79,11 @@ def session_close_tracker():
def patched_external_dependencies():
"""Patch non-DB collaborators while keeping database behavior real."""
with (
patch("tasks.document_indexing_task.IndexingRunner") as mock_indexing_runner,
patch("tasks.document_indexing_task.FeatureService") as mock_feature_service,
patch("tasks.document_indexing_task.generate_summary_index_task") as mock_summary_task,
patch("tasks.document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner,
patch("tasks.document_indexing_task.FeatureService", autospec=True) as mock_feature_service,
patch("tasks.document_indexing_task.generate_summary_index_task", autospec=True) as mock_summary_task,
):
mock_runner_instance = MagicMock()
mock_indexing_runner.return_value = mock_runner_instance
mock_runner_instance = mock_indexing_runner.return_value
mock_features = MagicMock()
mock_features.billing.enabled = False
mock_features.billing.subscription.plan = CloudPlan.PROFESSIONAL
@ -307,9 +307,17 @@ class TestDatasetIndexingTaskIntegration:
# Act
with (
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[next_task]),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time") as set_waiting_spy,
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key") as delete_key_spy,
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks",
return_value=[next_task],
autospec=True,
),
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time", autospec=True
) as set_waiting_spy,
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key", autospec=True
) as delete_key_spy,
):
_document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
@ -336,8 +344,10 @@ class TestDatasetIndexingTaskIntegration:
# Act
with (
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[]),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key") as delete_key_spy,
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[], autospec=True),
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key", autospec=True
) as delete_key_spy,
):
_document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
@ -426,9 +436,13 @@ class TestDatasetIndexingTaskIntegration:
# Act
with (
patch("tasks.document_indexing_task._document_indexing", side_effect=Exception("failed")),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[next_task]),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time"),
patch("tasks.document_indexing_task._document_indexing", side_effect=Exception("failed"), autospec=True),
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks",
return_value=[next_task],
autospec=True,
),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time", autospec=True),
):
_document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
@ -511,8 +525,11 @@ class TestDatasetIndexingTaskIntegration:
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks",
return_value=pending_tasks[:concurrency_limit],
autospec=True,
),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time") as set_waiting_spy,
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time", autospec=True
) as set_waiting_spy,
):
_document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
@ -538,8 +555,12 @@ class TestDatasetIndexingTaskIntegration:
# Act
with (
patch("tasks.document_indexing_task.dify_config.TENANT_ISOLATED_TASK_CONCURRENCY", 3),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=ordered_tasks),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time"),
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks",
return_value=ordered_tasks,
autospec=True,
),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time", autospec=True),
):
_document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
@ -578,8 +599,10 @@ class TestDatasetIndexingTaskIntegration:
# Act
with (
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[]),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key") as delete_key_spy,
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[], autospec=True),
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key", autospec=True
) as delete_key_spy,
):
normal_document_indexing_task(dataset.tenant_id, dataset.id, document_ids)
@ -599,8 +622,10 @@ class TestDatasetIndexingTaskIntegration:
# Act
with (
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[]),
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key") as delete_key_spy,
patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[], autospec=True),
patch(
"tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key", autospec=True
) as delete_key_spy,
):
priority_document_indexing_task(dataset.tenant_id, dataset.id, document_ids)

View File

@ -216,7 +216,7 @@ class TestDeleteSegmentFromIndexTask:
db_session_with_containers.commit()
return segments
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
def test_delete_segment_from_index_task_success(self, mock_index_processor_factory, db_session_with_containers):
"""
Test successful segment deletion from index with comprehensive verification.
@ -399,7 +399,7 @@ class TestDeleteSegmentFromIndexTask:
# Verify the task completed without exceptions
assert result is None # Task should return None when indexing is not completed
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
def test_delete_segment_from_index_task_index_processor_clean(
self, mock_index_processor_factory, db_session_with_containers
):
@ -457,7 +457,7 @@ class TestDeleteSegmentFromIndexTask:
mock_index_processor_factory.reset_mock()
mock_processor.reset_mock()
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
def test_delete_segment_from_index_task_exception_handling(
self, mock_index_processor_factory, db_session_with_containers
):
@ -501,7 +501,7 @@ class TestDeleteSegmentFromIndexTask:
assert call_args[1]["with_keywords"] is True
assert call_args[1]["delete_child_chunks"] is True
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
def test_delete_segment_from_index_task_empty_index_node_ids(
self, mock_index_processor_factory, db_session_with_containers
):
@ -543,7 +543,7 @@ class TestDeleteSegmentFromIndexTask:
assert call_args[1]["with_keywords"] is True
assert call_args[1]["delete_child_chunks"] is True
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
def test_delete_segment_from_index_task_large_index_node_ids(
self, mock_index_processor_factory, db_session_with_containers
):

View File

@ -32,14 +32,11 @@ class TestDocumentIndexingTasks:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.document_indexing_task.IndexingRunner") as mock_indexing_runner,
patch("tasks.document_indexing_task.FeatureService") as mock_feature_service,
patch("tasks.document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner,
patch("tasks.document_indexing_task.FeatureService", autospec=True) as mock_feature_service,
):
# Setup mock indexing runner
mock_runner_instance = MagicMock()
mock_indexing_runner.return_value = mock_runner_instance
# Setup mock feature service
mock_runner_instance = mock_indexing_runner.return_value # Setup mock feature service
mock_features = MagicMock()
mock_features.billing.enabled = False
mock_feature_service.get_features.return_value = mock_features

View File

@ -16,15 +16,13 @@ class TestDocumentIndexingUpdateTask:
- IndexingRunner.run([...])
"""
with (
patch("tasks.document_indexing_update_task.IndexProcessorFactory") as mock_factory,
patch("tasks.document_indexing_update_task.IndexingRunner") as mock_runner,
patch("tasks.document_indexing_update_task.IndexProcessorFactory", autospec=True) as mock_factory,
patch("tasks.document_indexing_update_task.IndexingRunner", autospec=True) as mock_runner,
):
processor_instance = MagicMock()
mock_factory.return_value.init_index_processor.return_value = processor_instance
runner_instance = MagicMock()
mock_runner.return_value = runner_instance
runner_instance = mock_runner.return_value
yield {
"factory": mock_factory,
"processor": processor_instance,

View File

@ -31,15 +31,14 @@ class TestDuplicateDocumentIndexingTasks:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.duplicate_document_indexing_task.IndexingRunner") as mock_indexing_runner,
patch("tasks.duplicate_document_indexing_task.FeatureService") as mock_feature_service,
patch("tasks.duplicate_document_indexing_task.IndexProcessorFactory") as mock_index_processor_factory,
patch("tasks.duplicate_document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner,
patch("tasks.duplicate_document_indexing_task.FeatureService", autospec=True) as mock_feature_service,
patch(
"tasks.duplicate_document_indexing_task.IndexProcessorFactory", autospec=True
) as mock_index_processor_factory,
):
# Setup mock indexing runner
mock_runner_instance = MagicMock()
mock_indexing_runner.return_value = mock_runner_instance
# Setup mock feature service
mock_runner_instance = mock_indexing_runner.return_value # Setup mock feature service
mock_features = MagicMock()
mock_features.billing.enabled = False
mock_feature_service.get_features.return_value = mock_features
@ -650,7 +649,7 @@ class TestDuplicateDocumentIndexingTasks:
updated_document = db_session_with_containers.query(Document).where(Document.id == doc_id).first()
assert updated_document.indexing_status == "parsing"
@patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue")
@patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue", autospec=True)
def test_normal_duplicate_document_indexing_task_with_tenant_queue(
self, mock_queue_class, db_session_with_containers, mock_external_service_dependencies
):
@ -693,7 +692,7 @@ class TestDuplicateDocumentIndexingTasks:
updated_document = db_session_with_containers.query(Document).where(Document.id == doc_id).first()
assert updated_document.indexing_status == "parsing"
@patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue")
@patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue", autospec=True)
def test_priority_duplicate_document_indexing_task_with_tenant_queue(
self, mock_queue_class, db_session_with_containers, mock_external_service_dependencies
):
@ -737,7 +736,7 @@ class TestDuplicateDocumentIndexingTasks:
updated_document = db_session_with_containers.query(Document).where(Document.id == doc_id).first()
assert updated_document.indexing_status == "parsing"
@patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue")
@patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue", autospec=True)
def test_tenant_queue_wrapper_processes_next_tasks(
self, mock_queue_class, db_session_with_containers, mock_external_service_dependencies
):

View File

@ -18,7 +18,9 @@ class TestEnableSegmentsToIndexTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.enable_segments_to_index_task.IndexProcessorFactory") as mock_index_processor_factory,
patch(
"tasks.enable_segments_to_index_task.IndexProcessorFactory", autospec=True
) as mock_index_processor_factory,
):
# Setup mock index processor
mock_processor = MagicMock()
@ -370,7 +372,7 @@ class TestEnableSegmentsToIndexTask:
redis_client.set(indexing_cache_key, "processing", ex=300)
# Mock the get_child_chunks method for each segment
with patch.object(DocumentSegment, "get_child_chunks") as mock_get_child_chunks:
with patch.object(DocumentSegment, "get_child_chunks", autospec=True) as mock_get_child_chunks:
# Setup mock to return child chunks for each segment
mock_child_chunks = []
for i in range(2): # Each segment has 2 child chunks

View File

@ -1,4 +1,4 @@
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import pytest
from faker import Faker
@ -16,16 +16,14 @@ class TestMailAccountDeletionTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.mail_account_deletion_task.mail") as mock_mail,
patch("tasks.mail_account_deletion_task.get_email_i18n_service") as mock_get_email_service,
patch("tasks.mail_account_deletion_task.mail", autospec=True) as mock_mail,
patch("tasks.mail_account_deletion_task.get_email_i18n_service", autospec=True) as mock_get_email_service,
):
# Setup mock mail service
mock_mail.is_inited.return_value = True
# Setup mock email service
mock_email_service = MagicMock()
mock_get_email_service.return_value = mock_email_service
mock_email_service = mock_get_email_service.return_value
yield {
"mail": mock_mail,
"get_email_service": mock_get_email_service,

View File

@ -1,4 +1,4 @@
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import pytest
from faker import Faker
@ -15,16 +15,14 @@ class TestMailChangeMailTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.mail_change_mail_task.mail") as mock_mail,
patch("tasks.mail_change_mail_task.get_email_i18n_service") as mock_get_email_i18n_service,
patch("tasks.mail_change_mail_task.mail", autospec=True) as mock_mail,
patch("tasks.mail_change_mail_task.get_email_i18n_service", autospec=True) as mock_get_email_i18n_service,
):
# Setup mock mail service
mock_mail.is_inited.return_value = True
# Setup mock email i18n service
mock_email_service = MagicMock()
mock_get_email_i18n_service.return_value = mock_email_service
mock_email_service = mock_get_email_i18n_service.return_value
yield {
"mail": mock_mail,
"email_i18n_service": mock_email_service,

View File

@ -53,8 +53,8 @@ class TestSendEmailCodeLoginMailTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.mail_email_code_login.mail") as mock_mail,
patch("tasks.mail_email_code_login.get_email_i18n_service") as mock_email_service,
patch("tasks.mail_email_code_login.mail", autospec=True) as mock_mail,
patch("tasks.mail_email_code_login.get_email_i18n_service", autospec=True) as mock_email_service,
):
# Setup default mock returns
mock_mail.is_inited.return_value = True
@ -573,7 +573,7 @@ class TestSendEmailCodeLoginMailTask:
mock_email_service_instance.send_email.side_effect = exception
# Mock logging to capture error messages
with patch("tasks.mail_email_code_login.logger") as mock_logger:
with patch("tasks.mail_email_code_login.logger", autospec=True) as mock_logger:
# Act: Execute the task - it should handle the exception gracefully
send_email_code_login_mail_task(
language=test_language,

View File

@ -1,4 +1,4 @@
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import pytest
from faker import Faker
@ -13,18 +13,15 @@ class TestMailInnerTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.mail_inner_task.mail") as mock_mail,
patch("tasks.mail_inner_task.get_email_i18n_service") as mock_get_email_i18n_service,
patch("tasks.mail_inner_task._render_template_with_strategy") as mock_render_template,
patch("tasks.mail_inner_task.mail", autospec=True) as mock_mail,
patch("tasks.mail_inner_task.get_email_i18n_service", autospec=True) as mock_get_email_i18n_service,
patch("tasks.mail_inner_task._render_template_with_strategy", autospec=True) as mock_render_template,
):
# Setup mock mail service
mock_mail.is_inited.return_value = True
# Setup mock email i18n service
mock_email_service = MagicMock()
mock_get_email_i18n_service.return_value = mock_email_service
# Setup mock template rendering
mock_email_service = mock_get_email_i18n_service.return_value # Setup mock template rendering
mock_render_template.return_value = "<html>Test email content</html>"
yield {

View File

@ -56,9 +56,9 @@ class TestMailInviteMemberTask:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("tasks.mail_invite_member_task.mail") as mock_mail,
patch("tasks.mail_invite_member_task.get_email_i18n_service") as mock_email_service,
patch("tasks.mail_invite_member_task.dify_config") as mock_config,
patch("tasks.mail_invite_member_task.mail", autospec=True) as mock_mail,
patch("tasks.mail_invite_member_task.get_email_i18n_service", autospec=True) as mock_email_service,
patch("tasks.mail_invite_member_task.dify_config", autospec=True) as mock_config,
):
# Setup mail service mock
mock_mail.is_inited.return_value = True
@ -306,7 +306,7 @@ class TestMailInviteMemberTask:
mock_email_service.send_email.side_effect = Exception("Email service failed")
# Act & Assert: Execute task and verify exception is handled
with patch("tasks.mail_invite_member_task.logger") as mock_logger:
with patch("tasks.mail_invite_member_task.logger", autospec=True) as mock_logger:
send_invite_member_mail_task(
language="en-US",
to="test@example.com",

View File

@ -7,7 +7,7 @@ testing with actual database and service dependencies.
"""
import logging
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import pytest
from faker import Faker
@ -30,16 +30,14 @@ class TestMailOwnerTransferTask:
def mock_mail_dependencies(self):
"""Mock setup for mail service dependencies."""
with (
patch("tasks.mail_owner_transfer_task.mail") as mock_mail,
patch("tasks.mail_owner_transfer_task.get_email_i18n_service") as mock_get_email_service,
patch("tasks.mail_owner_transfer_task.mail", autospec=True) as mock_mail,
patch("tasks.mail_owner_transfer_task.get_email_i18n_service", autospec=True) as mock_get_email_service,
):
# Setup mock mail service
mock_mail.is_inited.return_value = True
# Setup mock email service
mock_email_service = MagicMock()
mock_get_email_service.return_value = mock_email_service
mock_email_service = mock_get_email_service.return_value
yield {
"mail": mock_mail,
"email_service": mock_email_service,

View File

@ -5,7 +5,7 @@ This module provides integration tests for email registration tasks
using TestContainers to ensure real database and service interactions.
"""
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import pytest
from faker import Faker
@ -21,16 +21,14 @@ class TestMailRegisterTask:
def mock_mail_dependencies(self):
"""Mock setup for mail service dependencies."""
with (
patch("tasks.mail_register_task.mail") as mock_mail,
patch("tasks.mail_register_task.get_email_i18n_service") as mock_get_email_service,
patch("tasks.mail_register_task.mail", autospec=True) as mock_mail,
patch("tasks.mail_register_task.get_email_i18n_service", autospec=True) as mock_get_email_service,
):
# Setup mock mail service
mock_mail.is_inited.return_value = True
# Setup mock email i18n service
mock_email_service = MagicMock()
mock_get_email_service.return_value = mock_email_service
mock_email_service = mock_get_email_service.return_value
yield {
"mail": mock_mail,
"email_service": mock_email_service,
@ -76,7 +74,7 @@ class TestMailRegisterTask:
to_email = fake.email()
code = fake.numerify("######")
with patch("tasks.mail_register_task.logger") as mock_logger:
with patch("tasks.mail_register_task.logger", autospec=True) as mock_logger:
send_email_register_mail_task(language="en-US", to=to_email, code=code)
mock_logger.exception.assert_called_once_with("Send email register mail to %s failed", to_email)
@ -89,7 +87,7 @@ class TestMailRegisterTask:
to_email = fake.email()
account_name = fake.name()
with patch("tasks.mail_register_task.dify_config") as mock_config:
with patch("tasks.mail_register_task.dify_config", autospec=True) as mock_config:
mock_config.CONSOLE_WEB_URL = "https://console.dify.ai"
send_email_register_mail_task_when_account_exist(language=language, to=to_email, account_name=account_name)
@ -129,6 +127,6 @@ class TestMailRegisterTask:
to_email = fake.email()
account_name = fake.name()
with patch("tasks.mail_register_task.logger") as mock_logger:
with patch("tasks.mail_register_task.logger", autospec=True) as mock_logger:
send_email_register_mail_task_when_account_exist(language="en-US", to=to_email, account_name=account_name)
mock_logger.exception.assert_called_once_with("Send email register mail to %s failed", to_email)