youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-03-04 15:26:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into feat/trigger

Browse Source
This commit is contained in:
yessenia
2025-09-25 17:14:24 +08:00
parent ef9a741781 407323f817
commit cd07eef639
3013 changed files with 148826 additions and 44294 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
api/tests/test_containers_integration_tests/conftest.py
View File

@ -11,7 +11,6 @@ import logging
import os
from collections.abc import Generator
from pathlib import Path
from typing import Optional
import pytest
from flask import Flask
@ -24,7 +23,7 @@ from testcontainers.postgres import PostgresContainer
from testcontainers.redis import RedisContainer
from app_factory import create_app
from models import db
from extensions.ext_database import db
# Configure logging for test containers
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
@ -42,14 +41,14 @@ class DifyTestContainers:
def __init__(self):
"""Initialize container management with default configurations."""
self.postgres: Optional[PostgresContainer] = None
self.redis: Optional[RedisContainer] = None
self.dify_sandbox: Optional[DockerContainer] = None
self.dify_plugin_daemon: Optional[DockerContainer] = None
self.postgres: PostgresContainer | None = None
self.redis: RedisContainer | None = None
self.dify_sandbox: DockerContainer | None = None
self.dify_plugin_daemon: DockerContainer | None = None
self._containers_started = False
logger.info("DifyTestContainers initialized - ready to manage test containers")
def start_containers_with_env(self) -> None:
def start_containers_with_env(self):
"""
Start all required containers for integration testing.
@ -174,7 +173,7 @@ class DifyTestContainers:
# Start Dify Plugin Daemon container for plugin management
# Dify Plugin Daemon provides plugin lifecycle management and execution
logger.info("Initializing Dify Plugin Daemon container...")
self.dify_plugin_daemon = DockerContainer(image="langgenius/dify-plugin-daemon:0.2.0-local")
self.dify_plugin_daemon = DockerContainer(image="langgenius/dify-plugin-daemon:0.3.0-local")
self.dify_plugin_daemon.with_exposed_ports(5002)
self.dify_plugin_daemon.env = {
"DB_HOST": db_host,
@ -230,7 +229,7 @@ class DifyTestContainers:
self._containers_started = True
logger.info("All test containers started successfully")
def stop_containers(self) -> None:
def stop_containers(self):
"""
Stop and clean up all test containers.
@ -345,6 +344,12 @@ def _create_app_with_containers() -> Flask:
with db.engine.connect() as conn, conn.begin():
conn.execute(text(_UUIDv7SQL))
db.create_all()
# migration_dir = _get_migration_dir()
# alembic_config = Config()
# alembic_config.config_file_name = str(migration_dir / "alembic.ini")
# alembic_config.set_main_option("sqlalchemy.url", _get_engine_url(db.engine))
# alembic_config.set_main_option("script_location", str(migration_dir))
# alembic_command.upgrade(revision="head", config=alembic_config)
logger.info("Database schema created successfully")
logger.info("Flask application configured and ready for testing")

28
api/tests/test_containers_integration_tests/factories/test_storage_key_loader.py
View File

@ -1,6 +1,5 @@
import unittest
from datetime import UTC, datetime
from typing import Optional
from unittest.mock import patch
from uuid import uuid4
@ -42,7 +41,7 @@ class TestStorageKeyLoader(unittest.TestCase):
self.session.rollback()
def _create_upload_file(
self, file_id: Optional[str] = None, storage_key: Optional[str] = None, tenant_id: Optional[str] = None
self, file_id: str | None = None, storage_key: str | None = None, tenant_id: str | None = None
) -> UploadFile:
"""Helper method to create an UploadFile record for testing."""
if file_id is None:
@ -74,7 +73,7 @@ class TestStorageKeyLoader(unittest.TestCase):
return upload_file
def _create_tool_file(
self, file_id: Optional[str] = None, file_key: Optional[str] = None, tenant_id: Optional[str] = None
self, file_id: str | None = None, file_key: str | None = None, tenant_id: str | None = None
) -> ToolFile:
"""Helper method to create a ToolFile record for testing."""
if file_id is None:
@ -84,16 +83,17 @@ class TestStorageKeyLoader(unittest.TestCase):
if tenant_id is None:
tenant_id = self.tenant_id
tool_file = ToolFile()
tool_file = ToolFile(
user_id=self.user_id,
tenant_id=tenant_id,
conversation_id=self.conversation_id,
file_key=file_key,
mimetype="text/plain",
original_url="http://example.com/file.txt",
name="test_tool_file.txt",
size=2048,
)
tool_file.id = file_id
tool_file.user_id = self.user_id
tool_file.tenant_id = tenant_id
tool_file.conversation_id = self.conversation_id
tool_file.file_key = file_key
tool_file.mimetype = "text/plain"
tool_file.original_url = "http://example.com/file.txt"
tool_file.name = "test_tool_file.txt"
tool_file.size = 2048
self.session.add(tool_file)
self.session.flush()
@ -101,9 +101,7 @@ class TestStorageKeyLoader(unittest.TestCase):
return tool_file
def _create_file(
self, related_id: str, transfer_method: FileTransferMethod, tenant_id: Optional[str] = None
) -> File:
def _create_file(self, related_id: str, transfer_method: FileTransferMethod, tenant_id: str | None = None) -> File:
"""Helper method to create a File object for testing."""
if tenant_id is None:
tenant_id = self.tenant_id

32
api/tests/test_containers_integration_tests/services/test_account_service.py
View File

@ -13,7 +13,6 @@ from services.account_service import AccountService, RegisterService, TenantServ
from services.errors.account import (
AccountAlreadyInTenantError,
AccountLoginError,
AccountNotFoundError,
AccountPasswordError,
AccountRegisterError,
CurrentPasswordIncorrectError,
@ -91,6 +90,28 @@ class TestAccountService:
assert account.password is None
assert account.password_salt is None
def test_create_account_password_invalid_new_password(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test account create with invalid new password format.
"""
fake = Faker()
email = fake.email()
name = fake.name()
# Setup mocks
mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
# Test with too short password (assuming minimum length validation)
with pytest.raises(ValueError): # Password validation error
AccountService.create_account(
email=email,
name=name,
interface_language="en-US",
password="invalid_new_password",
)
def test_create_account_registration_disabled(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test account creation when registration is disabled.
@ -139,7 +160,7 @@ class TestAccountService:
fake = Faker()
email = fake.email()
password = fake.password(length=12)
with pytest.raises(AccountNotFoundError):
with pytest.raises(AccountPasswordError):
AccountService.authenticate(email, password)
def test_authenticate_banned_account(self, db_session_with_containers, mock_external_service_dependencies):
@ -940,7 +961,8 @@ class TestAccountService:
Test getting user through non-existent email.
"""
fake = Faker()
non_existent_email = fake.email()
domain = f"test-{fake.random_letters(10)}.com"
non_existent_email = fake.email(domain=domain)
found_user = AccountService.get_user_through_email(non_existent_email)
assert found_user is None
@ -3278,7 +3300,7 @@ class TestRegisterService:
redis_client.setex(cache_key, 24 * 60 * 60, account_id)
# Execute invitation retrieval
result = RegisterService._get_invitation_by_token(
result = RegisterService.get_invitation_by_token(
token=token,
workspace_id=workspace_id,
email=email,
@ -3316,7 +3338,7 @@ class TestRegisterService:
redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
# Execute invitation retrieval
result = RegisterService._get_invitation_by_token(token=token)
result = RegisterService.get_invitation_by_token(token=token)
# Verify result contains expected data
assert result is not None

60
api/tests/test_containers_integration_tests/services/test_advanced_prompt_template_service.py
View File

@ -42,7 +42,7 @@ class TestAdvancedPromptTemplateService:
# Test data for Baichuan model
args = {
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "completion",
"model_name": "baichuan-13b-chat",
"has_context": "true",
@ -77,7 +77,7 @@ class TestAdvancedPromptTemplateService:
# Test data for common model
args = {
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "completion",
"model_name": "gpt-3.5-turbo",
"has_context": "true",
@ -116,7 +116,7 @@ class TestAdvancedPromptTemplateService:
for model_name in test_cases:
args = {
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "completion",
"model_name": model_name,
"has_context": "true",
@ -144,7 +144,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.CHAT.value, "completion", "true")
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.CHAT, "completion", "true")
# Assert: Verify the expected outcomes
assert result is not None
@ -173,7 +173,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.CHAT.value, "chat", "true")
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.CHAT, "chat", "true")
# Assert: Verify the expected outcomes
assert result is not None
@ -202,7 +202,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.COMPLETION.value, "completion", "true")
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.COMPLETION, "completion", "true")
# Assert: Verify the expected outcomes
assert result is not None
@ -230,7 +230,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.COMPLETION.value, "chat", "true")
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.COMPLETION, "chat", "true")
# Assert: Verify the expected outcomes
assert result is not None
@ -257,7 +257,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.CHAT.value, "completion", "false")
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.CHAT, "completion", "false")
# Assert: Verify the expected outcomes
assert result is not None
@ -303,7 +303,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.CHAT.value, "unsupported_mode", "true")
result = AdvancedPromptTemplateService.get_common_prompt(AppMode.CHAT, "unsupported_mode", "true")
# Assert: Verify empty dict is returned
assert result == {}
@ -442,7 +442,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.CHAT.value, "completion", "true")
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.CHAT, "completion", "true")
# Assert: Verify the expected outcomes
assert result is not None
@ -473,7 +473,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.CHAT.value, "chat", "true")
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.CHAT, "chat", "true")
# Assert: Verify the expected outcomes
assert result is not None
@ -502,7 +502,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.COMPLETION.value, "completion", "true")
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.COMPLETION, "completion", "true")
# Assert: Verify the expected outcomes
assert result is not None
@ -530,7 +530,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.COMPLETION.value, "chat", "true")
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.COMPLETION, "chat", "true")
# Assert: Verify the expected outcomes
assert result is not None
@ -557,7 +557,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.CHAT.value, "completion", "false")
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.CHAT, "completion", "false")
# Assert: Verify the expected outcomes
assert result is not None
@ -603,7 +603,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Act: Execute the method under test
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.CHAT.value, "unsupported_mode", "true")
result = AdvancedPromptTemplateService.get_baichuan_prompt(AppMode.CHAT, "unsupported_mode", "true")
# Assert: Verify empty dict is returned
assert result == {}
@ -621,7 +621,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Test all app modes
app_modes = [AppMode.CHAT.value, AppMode.COMPLETION.value]
app_modes = [AppMode.CHAT, AppMode.COMPLETION]
model_modes = ["completion", "chat"]
for app_mode in app_modes:
@ -653,7 +653,7 @@ class TestAdvancedPromptTemplateService:
fake = Faker()
# Test all app modes
app_modes = [AppMode.CHAT.value, AppMode.COMPLETION.value]
app_modes = [AppMode.CHAT, AppMode.COMPLETION]
model_modes = ["completion", "chat"]
for app_mode in app_modes:
@ -686,10 +686,10 @@ class TestAdvancedPromptTemplateService:
# Test edge cases
edge_cases = [
{"app_mode": "", "model_mode": "completion", "model_name": "gpt-3.5-turbo", "has_context": "true"},
{"app_mode": AppMode.CHAT.value, "model_mode": "", "model_name": "gpt-3.5-turbo", "has_context": "true"},
{"app_mode": AppMode.CHAT.value, "model_mode": "completion", "model_name": "", "has_context": "true"},
{"app_mode": AppMode.CHAT, "model_mode": "", "model_name": "gpt-3.5-turbo", "has_context": "true"},
{"app_mode": AppMode.CHAT, "model_mode": "completion", "model_name": "", "has_context": "true"},
{
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "completion",
"model_name": "gpt-3.5-turbo",
"has_context": "",
@ -723,7 +723,7 @@ class TestAdvancedPromptTemplateService:
# Test with context
args = {
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "completion",
"model_name": "gpt-3.5-turbo",
"has_context": "true",
@ -757,7 +757,7 @@ class TestAdvancedPromptTemplateService:
# Test with context
args = {
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "completion",
"model_name": "baichuan-13b-chat",
"has_context": "true",
@ -786,25 +786,25 @@ class TestAdvancedPromptTemplateService:
# Test different scenarios
test_scenarios = [
{
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "completion",
"model_name": "gpt-3.5-turbo",
"has_context": "true",
},
{
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "chat",
"model_name": "gpt-3.5-turbo",
"has_context": "true",
},
{
"app_mode": AppMode.COMPLETION.value,
"app_mode": AppMode.COMPLETION,
"model_mode": "completion",
"model_name": "gpt-3.5-turbo",
"has_context": "true",
},
{
"app_mode": AppMode.COMPLETION.value,
"app_mode": AppMode.COMPLETION,
"model_mode": "chat",
"model_name": "gpt-3.5-turbo",
"has_context": "true",
@ -843,25 +843,25 @@ class TestAdvancedPromptTemplateService:
# Test different scenarios
test_scenarios = [
{
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "completion",
"model_name": "baichuan-13b-chat",
"has_context": "true",
},
{
"app_mode": AppMode.CHAT.value,
"app_mode": AppMode.CHAT,
"model_mode": "chat",
"model_name": "baichuan-13b-chat",
"has_context": "true",
},
{
"app_mode": AppMode.COMPLETION.value,
"app_mode": AppMode.COMPLETION,
"model_mode": "completion",
"model_name": "baichuan-13b-chat",
"has_context": "true",
},
{
"app_mode": AppMode.COMPLETION.value,
"app_mode": AppMode.COMPLETION,
"model_mode": "chat",
"model_name": "baichuan-13b-chat",
"has_context": "true",

5
api/tests/test_containers_integration_tests/services/test_agent_service.py
View File

@ -1,10 +1,11 @@
import json
from unittest.mock import MagicMock, patch
from unittest.mock import MagicMock, create_autospec, patch
import pytest
from faker import Faker
from core.plugin.impl.exc import PluginDaemonClientSideError
from models.account import Account
from models.model import AppModelConfig, Conversation, EndUser, Message, MessageAgentThought
from services.account_service import AccountService, TenantService
from services.agent_service import AgentService
@ -21,7 +22,7 @@ class TestAgentService:
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.current_user") as mock_current_user,
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,

7
api/tests/test_containers_integration_tests/services/test_annotation_service.py
View File

@ -1,9 +1,10 @@
from unittest.mock import patch
from unittest.mock import create_autospec, patch
import pytest
from faker import Faker
from werkzeug.exceptions import NotFound
from models.account import Account
from models.model import MessageAnnotation
from services.annotation_service import AppAnnotationService
from services.app_service import AppService
@ -24,7 +25,9 @@ class TestAnnotationService:
patch("services.annotation_service.enable_annotation_reply_task") as mock_enable_task,
patch("services.annotation_service.disable_annotation_reply_task") as mock_disable_task,
patch("services.annotation_service.batch_import_annotations_task") as mock_batch_import_task,
patch("services.annotation_service.current_user") as mock_current_user,
patch(
"services.annotation_service.current_user", create_autospec(Account, instance=True)
) as mock_current_user,
):
# Setup default mock returns
mock_account_feature_service.get_features.return_value.billing.enabled = False

82
api/tests/test_containers_integration_tests/services/test_app_dsl_service.py
View File

@ -322,7 +322,87 @@ class TestAppDslService:
# Verify workflow service was called
mock_external_service_dependencies["workflow_service"].return_value.get_draft_workflow.assert_called_once_with(
app
app, None
)
def test_export_dsl_with_workflow_id_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful DSL export with specific workflow ID.
"""
fake = Faker()
app, account = self._create_test_app_and_account(db_session_with_containers, mock_external_service_dependencies)
# Update app to workflow mode
app.mode = "workflow"
db_session_with_containers.commit()
# Mock workflow service to return a workflow when specific workflow_id is provided
mock_workflow = MagicMock()
mock_workflow.to_dict.return_value = {
"graph": {"nodes": [{"id": "start", "type": "start", "data": {"type": "start"}}], "edges": []},
"features": {},
"environment_variables": [],
"conversation_variables": [],
}
# Mock the get_draft_workflow method to return different workflows based on workflow_id
def mock_get_draft_workflow(app_model, workflow_id=None):
if workflow_id == "specific-workflow-id":
return mock_workflow
return None
mock_external_service_dependencies[
"workflow_service"
].return_value.get_draft_workflow.side_effect = mock_get_draft_workflow
# Export DSL with specific workflow ID
exported_dsl = AppDslService.export_dsl(app, include_secret=False, workflow_id="specific-workflow-id")
# Parse exported YAML
exported_data = yaml.safe_load(exported_dsl)
# Verify exported data structure
assert exported_data["kind"] == "app"
assert exported_data["app"]["name"] == app.name
assert exported_data["app"]["mode"] == "workflow"
# Verify workflow was exported
assert "workflow" in exported_data
assert "graph" in exported_data["workflow"]
assert "nodes" in exported_data["workflow"]["graph"]
# Verify dependencies were exported
assert "dependencies" in exported_data
assert isinstance(exported_data["dependencies"], list)
# Verify workflow service was called with specific workflow ID
mock_external_service_dependencies["workflow_service"].return_value.get_draft_workflow.assert_called_once_with(
app, "specific-workflow-id"
)
def test_export_dsl_with_invalid_workflow_id_raises_error(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test that export_dsl raises error when invalid workflow ID is provided.
"""
fake = Faker()
app, account = self._create_test_app_and_account(db_session_with_containers, mock_external_service_dependencies)
# Update app to workflow mode
app.mode = "workflow"
db_session_with_containers.commit()
# Mock workflow service to return None when invalid workflow ID is provided
mock_external_service_dependencies["workflow_service"].return_value.get_draft_workflow.return_value = None
# Export DSL with invalid workflow ID should raise ValueError
with pytest.raises(ValueError, match="Missing draft workflow configuration, please check."):
AppDslService.export_dsl(app, include_secret=False, workflow_id="invalid-workflow-id")
# Verify workflow service was called with the invalid workflow ID
mock_external_service_dependencies["workflow_service"].return_value.get_draft_workflow.assert_called_once_with(
app, "invalid-workflow-id"
)
def test_check_dependencies_success(self, db_session_with_containers, mock_external_service_dependencies):

46
api/tests/test_containers_integration_tests/services/test_app_service.py
View File

@ -1,9 +1,10 @@
from unittest.mock import patch
from unittest.mock import create_autospec, patch
import pytest
from faker import Faker
from constants.model_template import default_app_templates
from models.account import Account
from models.model import App, Site
from services.account_service import AccountService, TenantService
from services.app_service import AppService
@ -161,8 +162,13 @@ class TestAppService:
app_service = AppService()
created_app = app_service.create_app(tenant.id, app_args, account)
# Get app using the service
retrieved_app = app_service.get_app(created_app)
# Get app using the service - needs current_user mock
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.id = account.id
mock_current_user.current_tenant_id = account.current_tenant_id
with patch("services.app_service.current_user", mock_current_user):
retrieved_app = app_service.get_app(created_app)
# Verify retrieved app matches created app
assert retrieved_app.id == created_app.id
@ -406,7 +412,11 @@ class TestAppService:
"use_icon_as_answer_icon": True,
}
with patch("flask_login.utils._get_user", return_value=account):
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.id = account.id
mock_current_user.current_tenant_id = account.current_tenant_id
with patch("services.app_service.current_user", mock_current_user):
updated_app = app_service.update_app(app, update_args)
# Verify updated fields
@ -456,7 +466,11 @@ class TestAppService:
# Update app name
new_name = "New App Name"
with patch("flask_login.utils._get_user", return_value=account):
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.id = account.id
mock_current_user.current_tenant_id = account.current_tenant_id
with patch("services.app_service.current_user", mock_current_user):
updated_app = app_service.update_app_name(app, new_name)
assert updated_app.name == new_name
@ -504,7 +518,11 @@ class TestAppService:
# Update app icon
new_icon = "🌟"
new_icon_background = "#FFD93D"
with patch("flask_login.utils._get_user", return_value=account):
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.id = account.id
mock_current_user.current_tenant_id = account.current_tenant_id
with patch("services.app_service.current_user", mock_current_user):
updated_app = app_service.update_app_icon(app, new_icon, new_icon_background)
assert updated_app.icon == new_icon
@ -551,13 +569,17 @@ class TestAppService:
original_site_status = app.enable_site
# Update site status to disabled
with patch("flask_login.utils._get_user", return_value=account):
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.id = account.id
mock_current_user.current_tenant_id = account.current_tenant_id
with patch("services.app_service.current_user", mock_current_user):
updated_app = app_service.update_app_site_status(app, False)
assert updated_app.enable_site is False
assert updated_app.updated_by == account.id
# Update site status back to enabled
with patch("flask_login.utils._get_user", return_value=account):
with patch("services.app_service.current_user", mock_current_user):
updated_app = app_service.update_app_site_status(updated_app, True)
assert updated_app.enable_site is True
assert updated_app.updated_by == account.id
@ -602,13 +624,17 @@ class TestAppService:
original_api_status = app.enable_api
# Update API status to disabled
with patch("flask_login.utils._get_user", return_value=account):
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.id = account.id
mock_current_user.current_tenant_id = account.current_tenant_id
with patch("services.app_service.current_user", mock_current_user):
updated_app = app_service.update_app_api_status(app, False)
assert updated_app.enable_api is False
assert updated_app.updated_by == account.id
# Update API status back to enabled
with patch("flask_login.utils._get_user", return_value=account):
with patch("services.app_service.current_user", mock_current_user):
updated_app = app_service.update_app_api_status(updated_app, True)
assert updated_app.enable_api is True
assert updated_app.updated_by == account.id

240
api/tests/test_containers_integration_tests/services/test_file_service.py
View File

@ -1,9 +1,10 @@
import hashlib
from io import BytesIO
from unittest.mock import patch
from unittest.mock import create_autospec, patch
import pytest
from faker import Faker
from sqlalchemy import Engine
from werkzeug.exceptions import NotFound
from configs import dify_config
@ -17,6 +18,12 @@ from services.file_service import FileService
class TestFileService:
"""Integration tests for FileService using testcontainers."""
@pytest.fixture
def engine(self, db_session_with_containers):
bind = db_session_with_containers.get_bind()
assert isinstance(bind, Engine)
return bind
@pytest.fixture
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
@ -156,7 +163,7 @@ class TestFileService:
return upload_file
# Test upload_file method
def test_upload_file_success(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_file_success(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test successful file upload with valid parameters.
"""
@ -167,7 +174,7 @@ class TestFileService:
content = b"test file content"
mimetype = "application/pdf"
upload_file = FileService.upload_file(
upload_file = FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -187,13 +194,9 @@ class TestFileService:
# Verify storage was called
mock_external_service_dependencies["storage"].save.assert_called_once()
# Verify database state
from extensions.ext_database import db
db.session.refresh(upload_file)
assert upload_file.id is not None
def test_upload_file_with_end_user(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_file_with_end_user(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test file upload with end user instead of account.
"""
@ -204,7 +207,7 @@ class TestFileService:
content = b"test image content"
mimetype = "image/jpeg"
upload_file = FileService.upload_file(
upload_file = FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -215,7 +218,9 @@ class TestFileService:
assert upload_file.created_by == end_user.id
assert upload_file.created_by_role == CreatorUserRole.END_USER.value
def test_upload_file_with_datasets_source(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_file_with_datasets_source(
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file upload with datasets source parameter.
"""
@ -226,7 +231,7 @@ class TestFileService:
content = b"test file content"
mimetype = "application/pdf"
upload_file = FileService.upload_file(
upload_file = FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -239,7 +244,7 @@ class TestFileService:
assert upload_file.source_url == "https://example.com/source"
def test_upload_file_invalid_filename_characters(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file upload with invalid filename characters.
@ -252,14 +257,16 @@ class TestFileService:
mimetype = "text/plain"
with pytest.raises(ValueError, match="Filename contains invalid characters"):
FileService.upload_file(
FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
user=account,
)
def test_upload_file_filename_too_long(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_file_filename_too_long(
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file upload with filename that exceeds length limit.
"""
@ -272,7 +279,7 @@ class TestFileService:
content = b"test content"
mimetype = "text/plain"
upload_file = FileService.upload_file(
upload_file = FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -288,7 +295,7 @@ class TestFileService:
assert len(base_name) <= 200
def test_upload_file_datasets_unsupported_type(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file upload for datasets with unsupported file type.
@ -301,7 +308,7 @@ class TestFileService:
mimetype = "image/jpeg"
with pytest.raises(UnsupportedFileTypeError):
FileService.upload_file(
FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -309,7 +316,7 @@ class TestFileService:
source="datasets",
)
def test_upload_file_too_large(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_file_too_large(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test file upload with file size exceeding limit.
"""
@ -322,7 +329,7 @@ class TestFileService:
mimetype = "image/jpeg"
with pytest.raises(FileTooLargeError):
FileService.upload_file(
FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -331,7 +338,7 @@ class TestFileService:
# Test is_file_size_within_limit method
def test_is_file_size_within_limit_image_success(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file size check for image files within limit.
@ -339,12 +346,12 @@ class TestFileService:
extension = "jpg"
file_size = dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 1024 # Exactly at limit
result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
result = FileService(engine).is_file_size_within_limit(extension=extension, file_size=file_size)
assert result is True
def test_is_file_size_within_limit_video_success(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file size check for video files within limit.
@ -352,12 +359,12 @@ class TestFileService:
extension = "mp4"
file_size = dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT * 1024 * 1024 # Exactly at limit
result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
result = FileService(engine).is_file_size_within_limit(extension=extension, file_size=file_size)
assert result is True
def test_is_file_size_within_limit_audio_success(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file size check for audio files within limit.
@ -365,12 +372,12 @@ class TestFileService:
extension = "mp3"
file_size = dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT * 1024 * 1024 # Exactly at limit
result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
result = FileService(engine).is_file_size_within_limit(extension=extension, file_size=file_size)
assert result is True
def test_is_file_size_within_limit_document_success(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file size check for document files within limit.
@ -378,12 +385,12 @@ class TestFileService:
extension = "pdf"
file_size = dify_config.UPLOAD_FILE_SIZE_LIMIT * 1024 * 1024 # Exactly at limit
result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
result = FileService(engine).is_file_size_within_limit(extension=extension, file_size=file_size)
assert result is True
def test_is_file_size_within_limit_image_exceeded(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file size check for image files exceeding limit.
@ -391,12 +398,12 @@ class TestFileService:
extension = "jpg"
file_size = dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 1024 + 1 # Exceeds limit
result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
result = FileService(engine).is_file_size_within_limit(extension=extension, file_size=file_size)
assert result is False
def test_is_file_size_within_limit_unknown_extension(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file size check for unknown file extension.
@ -404,12 +411,12 @@ class TestFileService:
extension = "xyz"
file_size = dify_config.UPLOAD_FILE_SIZE_LIMIT * 1024 * 1024 # Uses default limit
result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
result = FileService(engine).is_file_size_within_limit(extension=extension, file_size=file_size)
assert result is True
# Test upload_text method
def test_upload_text_success(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_text_success(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test successful text upload.
"""
@ -417,25 +424,30 @@ class TestFileService:
text = "This is a test text content"
text_name = "test_text.txt"
# Mock current_user
with patch("services.file_service.current_user") as mock_current_user:
mock_current_user.current_tenant_id = str(fake.uuid4())
mock_current_user.id = str(fake.uuid4())
# Mock current_user using create_autospec
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.current_tenant_id = str(fake.uuid4())
mock_current_user.id = str(fake.uuid4())
upload_file = FileService.upload_text(text=text, text_name=text_name)
upload_file = FileService(engine).upload_text(
text=text,
text_name=text_name,
user_id=mock_current_user.id,
tenant_id=mock_current_user.current_tenant_id,
)
assert upload_file is not None
assert upload_file.name == text_name
assert upload_file.size == len(text)
assert upload_file.extension == "txt"
assert upload_file.mime_type == "text/plain"
assert upload_file.used is True
assert upload_file.used_by == mock_current_user.id
assert upload_file is not None
assert upload_file.name == text_name
assert upload_file.size == len(text)
assert upload_file.extension == "txt"
assert upload_file.mime_type == "text/plain"
assert upload_file.used is True
assert upload_file.used_by == mock_current_user.id
# Verify storage was called
mock_external_service_dependencies["storage"].save.assert_called_once()
# Verify storage was called
mock_external_service_dependencies["storage"].save.assert_called_once()
def test_upload_text_name_too_long(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_text_name_too_long(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test text upload with name that exceeds length limit.
"""
@ -443,19 +455,24 @@ class TestFileService:
text = "test content"
long_name = "a" * 250 # Longer than 200 characters
# Mock current_user
with patch("services.file_service.current_user") as mock_current_user:
mock_current_user.current_tenant_id = str(fake.uuid4())
mock_current_user.id = str(fake.uuid4())
# Mock current_user using create_autospec
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.current_tenant_id = str(fake.uuid4())
mock_current_user.id = str(fake.uuid4())
upload_file = FileService.upload_text(text=text, text_name=long_name)
upload_file = FileService(engine).upload_text(
text=text,
text_name=long_name,
user_id=mock_current_user.id,
tenant_id=mock_current_user.current_tenant_id,
)
# Verify name was truncated
assert len(upload_file.name) <= 200
assert upload_file.name == "a" * 200
# Verify name was truncated
assert len(upload_file.name) <= 200
assert upload_file.name == "a" * 200
# Test get_file_preview method
def test_get_file_preview_success(self, db_session_with_containers, mock_external_service_dependencies):
def test_get_file_preview_success(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test successful file preview generation.
"""
@ -471,12 +488,14 @@ class TestFileService:
db.session.commit()
result = FileService.get_file_preview(file_id=upload_file.id)
result = FileService(engine).get_file_preview(file_id=upload_file.id)
assert result == "extracted text content"
mock_external_service_dependencies["extract_processor"].load_from_upload_file.assert_called_once()
def test_get_file_preview_file_not_found(self, db_session_with_containers, mock_external_service_dependencies):
def test_get_file_preview_file_not_found(
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file preview with non-existent file.
"""
@ -484,10 +503,10 @@ class TestFileService:
non_existent_id = str(fake.uuid4())
with pytest.raises(NotFound, match="File not found"):
FileService.get_file_preview(file_id=non_existent_id)
FileService(engine).get_file_preview(file_id=non_existent_id)
def test_get_file_preview_unsupported_file_type(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file preview with unsupported file type.
@ -505,9 +524,11 @@ class TestFileService:
db.session.commit()
with pytest.raises(UnsupportedFileTypeError):
FileService.get_file_preview(file_id=upload_file.id)
FileService(engine).get_file_preview(file_id=upload_file.id)
def test_get_file_preview_text_truncation(self, db_session_with_containers, mock_external_service_dependencies):
def test_get_file_preview_text_truncation(
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file preview with text that exceeds preview limit.
"""
@ -527,13 +548,13 @@ class TestFileService:
long_text = "x" * 5000 # Longer than PREVIEW_WORDS_LIMIT
mock_external_service_dependencies["extract_processor"].load_from_upload_file.return_value = long_text
result = FileService.get_file_preview(file_id=upload_file.id)
result = FileService(engine).get_file_preview(file_id=upload_file.id)
assert len(result) == 3000 # PREVIEW_WORDS_LIMIT
assert result == "x" * 3000
# Test get_image_preview method
def test_get_image_preview_success(self, db_session_with_containers, mock_external_service_dependencies):
def test_get_image_preview_success(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test successful image preview generation.
"""
@ -553,7 +574,7 @@ class TestFileService:
nonce = "test_nonce"
sign = "test_signature"
generator, mime_type = FileService.get_image_preview(
generator, mime_type = FileService(engine).get_image_preview(
file_id=upload_file.id,
timestamp=timestamp,
nonce=nonce,
@ -564,7 +585,9 @@ class TestFileService:
assert mime_type == upload_file.mime_type
mock_external_service_dependencies["file_helpers"].verify_image_signature.assert_called_once()
def test_get_image_preview_invalid_signature(self, db_session_with_containers, mock_external_service_dependencies):
def test_get_image_preview_invalid_signature(
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test image preview with invalid signature.
"""
@ -582,14 +605,16 @@ class TestFileService:
sign = "invalid_signature"
with pytest.raises(NotFound, match="File not found or signature is invalid"):
FileService.get_image_preview(
FileService(engine).get_image_preview(
file_id=upload_file.id,
timestamp=timestamp,
nonce=nonce,
sign=sign,
)
def test_get_image_preview_file_not_found(self, db_session_with_containers, mock_external_service_dependencies):
def test_get_image_preview_file_not_found(
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test image preview with non-existent file.
"""
@ -601,7 +626,7 @@ class TestFileService:
sign = "test_signature"
with pytest.raises(NotFound, match="File not found or signature is invalid"):
FileService.get_image_preview(
FileService(engine).get_image_preview(
file_id=non_existent_id,
timestamp=timestamp,
nonce=nonce,
@ -609,7 +634,7 @@ class TestFileService:
)
def test_get_image_preview_unsupported_file_type(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test image preview with non-image file type.
@ -631,7 +656,7 @@ class TestFileService:
sign = "test_signature"
with pytest.raises(UnsupportedFileTypeError):
FileService.get_image_preview(
FileService(engine).get_image_preview(
file_id=upload_file.id,
timestamp=timestamp,
nonce=nonce,
@ -640,7 +665,7 @@ class TestFileService:
# Test get_file_generator_by_file_id method
def test_get_file_generator_by_file_id_success(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test successful file generator retrieval.
@ -655,7 +680,7 @@ class TestFileService:
nonce = "test_nonce"
sign = "test_signature"
generator, file_obj = FileService.get_file_generator_by_file_id(
generator, file_obj = FileService(engine).get_file_generator_by_file_id(
file_id=upload_file.id,
timestamp=timestamp,
nonce=nonce,
@ -663,11 +688,11 @@ class TestFileService:
)
assert generator is not None
assert file_obj == upload_file
assert file_obj.id == upload_file.id
mock_external_service_dependencies["file_helpers"].verify_file_signature.assert_called_once()
def test_get_file_generator_by_file_id_invalid_signature(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file generator retrieval with invalid signature.
@ -686,7 +711,7 @@ class TestFileService:
sign = "invalid_signature"
with pytest.raises(NotFound, match="File not found or signature is invalid"):
FileService.get_file_generator_by_file_id(
FileService(engine).get_file_generator_by_file_id(
file_id=upload_file.id,
timestamp=timestamp,
nonce=nonce,
@ -694,7 +719,7 @@ class TestFileService:
)
def test_get_file_generator_by_file_id_file_not_found(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file generator retrieval with non-existent file.
@ -707,7 +732,7 @@ class TestFileService:
sign = "test_signature"
with pytest.raises(NotFound, match="File not found or signature is invalid"):
FileService.get_file_generator_by_file_id(
FileService(engine).get_file_generator_by_file_id(
file_id=non_existent_id,
timestamp=timestamp,
nonce=nonce,
@ -715,7 +740,9 @@ class TestFileService:
)
# Test get_public_image_preview method
def test_get_public_image_preview_success(self, db_session_with_containers, mock_external_service_dependencies):
def test_get_public_image_preview_success(
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test successful public image preview generation.
"""
@ -731,14 +758,14 @@ class TestFileService:
db.session.commit()
generator, mime_type = FileService.get_public_image_preview(file_id=upload_file.id)
generator, mime_type = FileService(engine).get_public_image_preview(file_id=upload_file.id)
assert generator is not None
assert mime_type == upload_file.mime_type
mock_external_service_dependencies["storage"].load.assert_called_once()
def test_get_public_image_preview_file_not_found(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test public image preview with non-existent file.
@ -747,10 +774,10 @@ class TestFileService:
non_existent_id = str(fake.uuid4())
with pytest.raises(NotFound, match="File not found or signature is invalid"):
FileService.get_public_image_preview(file_id=non_existent_id)
FileService(engine).get_public_image_preview(file_id=non_existent_id)
def test_get_public_image_preview_unsupported_file_type(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test public image preview with non-image file type.
@ -768,10 +795,10 @@ class TestFileService:
db.session.commit()
with pytest.raises(UnsupportedFileTypeError):
FileService.get_public_image_preview(file_id=upload_file.id)
FileService(engine).get_public_image_preview(file_id=upload_file.id)
# Test edge cases and boundary conditions
def test_upload_file_empty_content(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_file_empty_content(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test file upload with empty content.
"""
@ -782,7 +809,7 @@ class TestFileService:
content = b""
mimetype = "text/plain"
upload_file = FileService.upload_file(
upload_file = FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -793,7 +820,7 @@ class TestFileService:
assert upload_file.size == 0
def test_upload_file_special_characters_in_name(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file upload with special characters in filename (but valid ones).
@ -805,7 +832,7 @@ class TestFileService:
content = b"test content"
mimetype = "text/plain"
upload_file = FileService.upload_file(
upload_file = FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -816,7 +843,7 @@ class TestFileService:
assert upload_file.name == filename
def test_upload_file_different_case_extensions(
self, db_session_with_containers, mock_external_service_dependencies
self, db_session_with_containers, engine, mock_external_service_dependencies
):
"""
Test file upload with different case extensions.
@ -828,7 +855,7 @@ class TestFileService:
content = b"test content"
mimetype = "application/pdf"
upload_file = FileService.upload_file(
upload_file = FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -838,7 +865,7 @@ class TestFileService:
assert upload_file is not None
assert upload_file.extension == "pdf" # Should be converted to lowercase
def test_upload_text_empty_text(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_text_empty_text(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test text upload with empty text.
"""
@ -846,17 +873,22 @@ class TestFileService:
text = ""
text_name = "empty.txt"
# Mock current_user
with patch("services.file_service.current_user") as mock_current_user:
mock_current_user.current_tenant_id = str(fake.uuid4())
mock_current_user.id = str(fake.uuid4())
# Mock current_user using create_autospec
mock_current_user = create_autospec(Account, instance=True)
mock_current_user.current_tenant_id = str(fake.uuid4())
mock_current_user.id = str(fake.uuid4())
upload_file = FileService.upload_text(text=text, text_name=text_name)
upload_file = FileService(engine).upload_text(
text=text,
text_name=text_name,
user_id=mock_current_user.id,
tenant_id=mock_current_user.current_tenant_id,
)
assert upload_file is not None
assert upload_file.size == 0
assert upload_file is not None
assert upload_file.size == 0
def test_file_size_limits_edge_cases(self, db_session_with_containers, mock_external_service_dependencies):
def test_file_size_limits_edge_cases(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test file size limits with edge case values.
"""
@ -868,15 +900,15 @@ class TestFileService:
("pdf", dify_config.UPLOAD_FILE_SIZE_LIMIT),
]:
file_size = limit_config * 1024 * 1024
result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
result = FileService(engine).is_file_size_within_limit(extension=extension, file_size=file_size)
assert result is True
# Test one byte over limit
file_size = limit_config * 1024 * 1024 + 1
result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
result = FileService(engine).is_file_size_within_limit(extension=extension, file_size=file_size)
assert result is False
def test_upload_file_with_source_url(self, db_session_with_containers, mock_external_service_dependencies):
def test_upload_file_with_source_url(self, db_session_with_containers, engine, mock_external_service_dependencies):
"""
Test file upload with source URL that gets overridden by signed URL.
"""
@ -888,7 +920,7 @@ class TestFileService:
mimetype = "application/pdf"
source_url = "https://original-source.com/file.pdf"
upload_file = FileService.upload_file(
upload_file = FileService(engine).upload_file(
filename=filename,
content=content,
mimetype=mimetype,
@ -901,7 +933,7 @@ class TestFileService:
# The signed URL should only be set when source_url is empty
# Let's test that scenario
upload_file2 = FileService.upload_file(
upload_file2 = FileService(engine).upload_file(
filename="test2.pdf",
content=b"test content 2",
mimetype="application/pdf",

30
api/tests/test_containers_integration_tests/services/test_metadata_service.py
View File

@ -1,4 +1,4 @@
from unittest.mock import patch
from unittest.mock import create_autospec, patch
import pytest
from faker import Faker
@ -17,7 +17,9 @@ class TestMetadataService:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.metadata_service.current_user") as mock_current_user,
patch(
"services.metadata_service.current_user", create_autospec(Account, instance=True)
) as mock_current_user,
patch("services.metadata_service.redis_client") as mock_redis_client,
patch("services.dataset_service.DocumentService") as mock_document_service,
):
@ -253,7 +255,7 @@ class TestMetadataService:
mock_external_service_dependencies["current_user"].id = account.id
# Try to create metadata with built-in field name
built_in_field_name = BuiltInField.document_name.value
built_in_field_name = BuiltInField.document_name
metadata_args = MetadataArgs(type="string", name=built_in_field_name)
# Act & Assert: Verify proper error handling
@ -373,7 +375,7 @@ class TestMetadataService:
metadata = MetadataService.create_metadata(dataset.id, metadata_args)
# Try to update with built-in field name
built_in_field_name = BuiltInField.document_name.value
built_in_field_name = BuiltInField.document_name
with pytest.raises(ValueError, match="Metadata name already exists in Built-in fields."):
MetadataService.update_metadata_name(dataset.id, metadata.id, built_in_field_name)
@ -538,11 +540,11 @@ class TestMetadataService:
field_names = [field["name"] for field in result]
field_types = [field["type"] for field in result]
assert BuiltInField.document_name.value in field_names
assert BuiltInField.uploader.value in field_names
assert BuiltInField.upload_date.value in field_names
assert BuiltInField.last_update_date.value in field_names
assert BuiltInField.source.value in field_names
assert BuiltInField.document_name in field_names
assert BuiltInField.uploader in field_names
assert BuiltInField.upload_date in field_names
assert BuiltInField.last_update_date in field_names
assert BuiltInField.source in field_names
# Verify field types
assert "string" in field_types
@ -680,11 +682,11 @@ class TestMetadataService:
# Set document metadata with built-in fields
document.doc_metadata = {
BuiltInField.document_name.value: document.name,
BuiltInField.uploader.value: "test_uploader",
BuiltInField.upload_date.value: 1234567890.0,
BuiltInField.last_update_date.value: 1234567890.0,
BuiltInField.source.value: "test_source",
BuiltInField.document_name: document.name,
BuiltInField.uploader: "test_uploader",
BuiltInField.upload_date: 1234567890.0,
BuiltInField.last_update_date: 1234567890.0,
BuiltInField.source: "test_source",
}
db.session.add(document)
db.session.commit()

7
api/tests/test_containers_integration_tests/services/test_model_load_balancing_service.py
View File

@ -2,6 +2,7 @@ from unittest.mock import MagicMock, patch
import pytest
from faker import Faker
from sqlalchemy import select
from models.account import TenantAccountJoin, TenantAccountRole
from models.model import Account, Tenant
@ -468,7 +469,7 @@ class TestModelLoadBalancingService:
assert load_balancing_config.id is not None
# Verify inherit config was created in database
inherit_configs = (
db.session.query(LoadBalancingModelConfig).where(LoadBalancingModelConfig.name == "__inherit__").all()
)
inherit_configs = db.session.scalars(
select(LoadBalancingModelConfig).where(LoadBalancingModelConfig.name == "__inherit__")
).all()
assert len(inherit_configs) == 1

13
api/tests/test_containers_integration_tests/services/test_tag_service.py
View File

@ -1,7 +1,8 @@
from unittest.mock import patch
from unittest.mock import create_autospec, patch
import pytest
from faker import Faker
from sqlalchemy import select
from werkzeug.exceptions import NotFound
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
@ -17,7 +18,7 @@ class TestTagService:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.tag_service.current_user") as mock_current_user,
patch("services.tag_service.current_user", create_autospec(Account, instance=True)) as mock_current_user,
):
# Setup default mock returns
mock_current_user.current_tenant_id = "test-tenant-id"
@ -954,7 +955,9 @@ class TestTagService:
from extensions.ext_database import db
# Verify only one binding exists
bindings = db.session.query(TagBinding).where(TagBinding.tag_id == tag.id, TagBinding.target_id == app.id).all()
bindings = db.session.scalars(
select(TagBinding).where(TagBinding.tag_id == tag.id, TagBinding.target_id == app.id)
).all()
assert len(bindings) == 1
def test_save_tag_binding_invalid_target_type(self, db_session_with_containers, mock_external_service_dependencies):
@ -1064,7 +1067,9 @@ class TestTagService:
# No error should be raised, and database state should remain unchanged
from extensions.ext_database import db
bindings = db.session.query(TagBinding).where(TagBinding.tag_id == tag.id, TagBinding.target_id == app.id).all()
bindings = db.session.scalars(
select(TagBinding).where(TagBinding.tag_id == tag.id, TagBinding.target_id == app.id)
).all()
assert len(bindings) == 0
def test_check_target_exists_knowledge_success(

9
api/tests/test_containers_integration_tests/services/test_web_conversation_service.py
View File

@ -2,6 +2,7 @@ from unittest.mock import patch
import pytest
from faker import Faker
from sqlalchemy import select
from core.app.entities.app_invoke_entities import InvokeFrom
from models.account import Account
@ -354,16 +355,14 @@ class TestWebConversationService:
# Verify only one pinned conversation record exists
from extensions.ext_database import db
pinned_conversations = (
db.session.query(PinnedConversation)
.where(
pinned_conversations = db.session.scalars(
select(PinnedConversation).where(
PinnedConversation.app_id == app.id,
PinnedConversation.conversation_id == conversation.id,
PinnedConversation.created_by_role == "account",
PinnedConversation.created_by == account.id,
)
.all()
)
).all()
assert len(pinned_conversations) == 1

35
api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py
View File

@ -1,3 +1,5 @@
import time
import uuid
from unittest.mock import patch
import pytest
@ -57,10 +59,12 @@ class TestWebAppAuthService:
tuple: (account, tenant) - Created account and tenant instances
"""
fake = Faker()
import uuid
# Create account
# Create account with unique email to avoid collisions
unique_email = f"test_{uuid.uuid4().hex[:8]}@example.com"
account = Account(
email=fake.email(),
email=unique_email,
name=fake.name(),
interface_language="en-US",
status="active",
@ -109,8 +113,11 @@ class TestWebAppAuthService:
password = fake.password(length=12)
# Create account with password
import uuid
unique_email = f"test_{uuid.uuid4().hex[:8]}@example.com"
account = Account(
email=fake.email(),
email=unique_email,
name=fake.name(),
interface_language="en-US",
status="active",
@ -243,9 +250,15 @@ class TestWebAppAuthService:
- Proper error handling for non-existent accounts
- Correct exception type and message
"""
# Arrange: Use non-existent email
fake = Faker()
non_existent_email = fake.email()
# Arrange: Generate a guaranteed non-existent email
# Use UUID and timestamp to ensure uniqueness
unique_id = str(uuid.uuid4()).replace("-", "")
timestamp = str(int(time.time() * 1000000)) # microseconds
non_existent_email = f"nonexistent_{unique_id}_{timestamp}@test-domain-that-never-exists.invalid"
# Double-check this email doesn't exist in the database
existing_account = db_session_with_containers.query(Account).filter_by(email=non_existent_email).first()
assert existing_account is None, f"Test email {non_existent_email} already exists in database"
# Act & Assert: Verify proper error handling
with pytest.raises(AccountNotFoundError):
@ -322,9 +335,12 @@ class TestWebAppAuthService:
"""
# Arrange: Create account without password
fake = Faker()
import uuid
unique_email = f"test_{uuid.uuid4().hex[:8]}@example.com"
account = Account(
email=fake.email(),
email=unique_email,
name=fake.name(),
interface_language="en-US",
status="active",
@ -431,9 +447,12 @@ class TestWebAppAuthService:
"""
# Arrange: Create banned account
fake = Faker()
import uuid
unique_email = f"test_{uuid.uuid4().hex[:8]}@example.com"
account = Account(
email=fake.email(),
email=unique_email,
name=fake.name(),
interface_language="en-US",
status=AccountStatus.BANNED.value,

1437
api/tests/test_containers_integration_tests/services/test_website_service.py
View File

File diff suppressed because it is too large Load Diff

1
api/tests/test_containers_integration_tests/services/test_workflow_draft_variable_service.py
View File

@ -108,6 +108,7 @@ class TestWorkflowDraftVariableService:
created_by=app.created_by,
environment_variables=[],
conversation_variables=[],
rag_pipeline_variables=[],
)
from extensions.ext_database import db

135
api/tests/test_containers_integration_tests/services/test_workflow_service.py
View File

@ -96,7 +96,7 @@ class TestWorkflowService:
app.tenant_id = fake.uuid4()
app.name = fake.company()
app.description = fake.text()
app.mode = AppMode.WORKFLOW.value
app.mode = AppMode.WORKFLOW
app.icon_type = "emoji"
app.icon = "🤖"
app.icon_background = "#FFEAD5"
@ -883,7 +883,7 @@ class TestWorkflowService:
# Create chat mode app
app = self._create_test_app(db_session_with_containers, fake)
app.mode = AppMode.CHAT.value
app.mode = AppMode.CHAT
# Create app model config (required for conversion)
from models.model import AppModelConfig
@ -926,7 +926,7 @@ class TestWorkflowService:
# Assert
assert result is not None
assert result.mode == AppMode.ADVANCED_CHAT.value # CHAT mode converts to ADVANCED_CHAT, not WORKFLOW
assert result.mode == AppMode.ADVANCED_CHAT # CHAT mode converts to ADVANCED_CHAT, not WORKFLOW
assert result.name == conversion_args["name"]
assert result.icon == conversion_args["icon"]
assert result.icon_type == conversion_args["icon_type"]
@ -945,7 +945,7 @@ class TestWorkflowService:
# Create completion mode app
app = self._create_test_app(db_session_with_containers, fake)
app.mode = AppMode.COMPLETION.value
app.mode = AppMode.COMPLETION
# Create app model config (required for conversion)
from models.model import AppModelConfig
@ -988,7 +988,7 @@ class TestWorkflowService:
# Assert
assert result is not None
assert result.mode == AppMode.WORKFLOW.value
assert result.mode == AppMode.WORKFLOW
assert result.name == conversion_args["name"]
assert result.icon == conversion_args["icon"]
assert result.icon_type == conversion_args["icon_type"]
@ -1007,7 +1007,7 @@ class TestWorkflowService:
# Create workflow mode app (already in workflow mode)
app = self._create_test_app(db_session_with_containers, fake)
app.mode = AppMode.WORKFLOW.value
app.mode = AppMode.WORKFLOW
from extensions.ext_database import db
@ -1030,7 +1030,7 @@ class TestWorkflowService:
# Arrange
fake = Faker()
app = self._create_test_app(db_session_with_containers, fake)
app.mode = AppMode.ADVANCED_CHAT.value
app.mode = AppMode.ADVANCED_CHAT
from extensions.ext_database import db
@ -1061,7 +1061,7 @@ class TestWorkflowService:
# Arrange
fake = Faker()
app = self._create_test_app(db_session_with_containers, fake)
app.mode = AppMode.WORKFLOW.value
app.mode = AppMode.WORKFLOW
from extensions.ext_database import db
@ -1421,16 +1421,19 @@ class TestWorkflowService:
# Mock successful node execution
def mock_successful_invoke():
from core.workflow.entities.node_entities import NodeRunResult
from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
from core.workflow.nodes.base.node import BaseNode
from core.workflow.nodes.event import RunCompletedEvent
import uuid
from datetime import datetime
from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
from core.workflow.graph_events import NodeRunSucceededEvent
from core.workflow.node_events import NodeRunResult
from core.workflow.nodes.base.node import Node
# Create mock node
mock_node = MagicMock(spec=BaseNode)
mock_node.type_ = "start" # Use valid NodeType
mock_node = MagicMock(spec=Node)
mock_node.node_type = NodeType.START
mock_node.title = "Test Node"
mock_node.continue_on_error = False
mock_node.error_strategy = None
# Create mock result with valid metadata
mock_result = NodeRunResult(
@ -1441,25 +1444,37 @@ class TestWorkflowService:
metadata={"total_tokens": 100}, # Use valid metadata field
)
# Create mock event
mock_event = RunCompletedEvent(run_result=mock_result)
# Create mock event with all required fields
mock_event = NodeRunSucceededEvent(
id=str(uuid.uuid4()),
node_id=node_id,
node_type=NodeType.START,
node_run_result=mock_result,
start_at=datetime.now(),
)
return mock_node, [mock_event]
# Return node and generator
def event_generator():
yield mock_event
return mock_node, event_generator()
workflow_service = WorkflowService()
# Act
result = workflow_service._handle_node_run_result(
result = workflow_service._handle_single_step_result(
invoke_node_fn=mock_successful_invoke, start_at=start_at, node_id=node_id
)
# Assert
assert result is not None
assert result.node_id == node_id
assert result.node_type == "start" # Should match the mock node type
from core.workflow.enums import NodeType
assert result.node_type == NodeType.START # Should match the mock node type
assert result.title == "Test Node"
# Import the enum for comparison
from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
from core.workflow.enums import WorkflowNodeExecutionStatus
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
assert result.inputs is not None
@ -1481,34 +1496,47 @@ class TestWorkflowService:
# Mock failed node execution
def mock_failed_invoke():
from core.workflow.entities.node_entities import NodeRunResult
from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
from core.workflow.nodes.base.node import BaseNode
from core.workflow.nodes.event import RunCompletedEvent
import uuid
from datetime import datetime
from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
from core.workflow.graph_events import NodeRunFailedEvent
from core.workflow.node_events import NodeRunResult
from core.workflow.nodes.base.node import Node
# Create mock node
mock_node = MagicMock(spec=BaseNode)
mock_node.type_ = "llm" # Use valid NodeType
mock_node = MagicMock(spec=Node)
mock_node.node_type = NodeType.LLM
mock_node.title = "Test Node"
mock_node.continue_on_error = False
mock_node.error_strategy = None
# Create mock failed result
mock_result = NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
inputs={"input1": "value1"},
error="Test error message",
error_type="TestError",
)
# Create mock event
mock_event = RunCompletedEvent(run_result=mock_result)
# Create mock event with all required fields
mock_event = NodeRunFailedEvent(
id=str(uuid.uuid4()),
node_id=node_id,
node_type=NodeType.LLM,
node_run_result=mock_result,
error="Test error message",
start_at=datetime.now(),
)
return mock_node, [mock_event]
# Return node and generator
def event_generator():
yield mock_event
return mock_node, event_generator()
workflow_service = WorkflowService()
# Act
result = workflow_service._handle_node_run_result(
result = workflow_service._handle_single_step_result(
invoke_node_fn=mock_failed_invoke, start_at=start_at, node_id=node_id
)
@ -1516,7 +1544,7 @@ class TestWorkflowService:
assert result is not None
assert result.node_id == node_id
# Import the enum for comparison
from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
from core.workflow.enums import WorkflowNodeExecutionStatus
assert result.status == WorkflowNodeExecutionStatus.FAILED
assert result.error is not None
@ -1537,17 +1565,18 @@ class TestWorkflowService:
# Mock node execution with continue_on_error
def mock_continue_on_error_invoke():
from core.workflow.entities.node_entities import NodeRunResult
from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
from core.workflow.nodes.base.node import BaseNode
from core.workflow.nodes.enums import ErrorStrategy
from core.workflow.nodes.event import RunCompletedEvent
import uuid
from datetime import datetime
from core.workflow.enums import ErrorStrategy, NodeType, WorkflowNodeExecutionStatus
from core.workflow.graph_events import NodeRunFailedEvent
from core.workflow.node_events import NodeRunResult
from core.workflow.nodes.base.node import Node
# Create mock node with continue_on_error
mock_node = MagicMock(spec=BaseNode)
mock_node.type_ = "tool" # Use valid NodeType
mock_node = MagicMock(spec=Node)
mock_node.node_type = NodeType.TOOL
mock_node.title = "Test Node"
mock_node.continue_on_error = True
mock_node.error_strategy = ErrorStrategy.DEFAULT_VALUE
mock_node.default_value_dict = {"default_output": "default_value"}
@ -1556,18 +1585,28 @@ class TestWorkflowService:
status=WorkflowNodeExecutionStatus.FAILED,
inputs={"input1": "value1"},
error="Test error message",
error_type="TestError",
)
# Create mock event
mock_event = RunCompletedEvent(run_result=mock_result)
# Create mock event with all required fields
mock_event = NodeRunFailedEvent(
id=str(uuid.uuid4()),
node_id=node_id,
node_type=NodeType.TOOL,
node_run_result=mock_result,
error="Test error message",
start_at=datetime.now(),
)
return mock_node, [mock_event]
# Return node and generator
def event_generator():
yield mock_event
return mock_node, event_generator()
workflow_service = WorkflowService()
# Act
result = workflow_service._handle_node_run_result(
result = workflow_service._handle_single_step_result(
invoke_node_fn=mock_continue_on_error_invoke, start_at=start_at, node_id=node_id
)
@ -1575,7 +1614,7 @@ class TestWorkflowService:
assert result is not None
assert result.node_id == node_id
# Import the enum for comparison
from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
from core.workflow.enums import WorkflowNodeExecutionStatus
assert result.status == WorkflowNodeExecutionStatus.EXCEPTION # Should be EXCEPTION, not FAILED
assert result.outputs is not None

39
api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py
View File

@ -706,7 +706,14 @@ class TestMCPToolManageService:
# Verify mock interactions
mock_mcp_client.assert_called_once_with(
"https://example.com/mcp", mcp_provider.id, tenant.id, authed=False, for_list=True
"https://example.com/mcp",
mcp_provider.id,
tenant.id,
authed=False,
for_list=True,
headers={},
timeout=30.0,
sse_read_timeout=300.0,
)
def test_list_mcp_tool_from_remote_server_auth_error(
@ -1181,6 +1188,11 @@ class TestMCPToolManageService:
db_session_with_containers, mock_external_service_dependencies
)
# Create MCP provider first
mcp_provider = self._create_test_mcp_provider(
db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id
)
# Mock MCPClient and its context manager
mock_tools = [
type("MockTool", (), {"model_dump": lambda self: {"name": "test_tool_1", "description": "Test tool 1"}})(),
@ -1194,7 +1206,7 @@ class TestMCPToolManageService:
# Act: Execute the method under test
result = MCPToolManageService._re_connect_mcp_provider(
"https://example.com/mcp", "test_provider_id", tenant.id
"https://example.com/mcp", mcp_provider.id, tenant.id
)
# Assert: Verify the expected outcomes
@ -1213,7 +1225,14 @@ class TestMCPToolManageService:
# Verify mock interactions
mock_mcp_client.assert_called_once_with(
"https://example.com/mcp", "test_provider_id", tenant.id, authed=False, for_list=True
"https://example.com/mcp",
mcp_provider.id,
tenant.id,
authed=False,
for_list=True,
headers={},
timeout=30.0,
sse_read_timeout=300.0,
)
def test_re_connect_mcp_provider_auth_error(self, db_session_with_containers, mock_external_service_dependencies):
@ -1231,6 +1250,11 @@ class TestMCPToolManageService:
db_session_with_containers, mock_external_service_dependencies
)
# Create MCP provider first
mcp_provider = self._create_test_mcp_provider(
db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id
)
# Mock MCPClient to raise authentication error
with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client:
from core.mcp.error import MCPAuthError
@ -1240,7 +1264,7 @@ class TestMCPToolManageService:
# Act: Execute the method under test
result = MCPToolManageService._re_connect_mcp_provider(
"https://example.com/mcp", "test_provider_id", tenant.id
"https://example.com/mcp", mcp_provider.id, tenant.id
)
# Assert: Verify the expected outcomes
@ -1265,6 +1289,11 @@ class TestMCPToolManageService:
db_session_with_containers, mock_external_service_dependencies
)
# Create MCP provider first
mcp_provider = self._create_test_mcp_provider(
db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id
)
# Mock MCPClient to raise connection error
with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client:
from core.mcp.error import MCPError
@ -1274,4 +1303,4 @@ class TestMCPToolManageService:
# Act & Assert: Verify proper error handling
with pytest.raises(ValueError, match="Failed to re-connect MCP server: Connection failed"):
MCPToolManageService._re_connect_mcp_provider("https://example.com/mcp", "test_provider_id", tenant.id)
MCPToolManageService._re_connect_mcp_provider("https://example.com/mcp", mcp_provider.id, tenant.id)

788
api/tests/test_containers_integration_tests/services/tools/test_tools_transform_service.py Normal file
View File

@ -0,0 +1,788 @@
from unittest.mock import Mock, patch
import pytest
from faker import Faker
from core.tools.entities.api_entities import ToolProviderApiEntity
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_entities import ToolProviderType
from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider
from services.tools.tools_transform_service import ToolTransformService
class TestToolTransformService:
"""Integration tests for ToolTransformService using testcontainers."""
@pytest.fixture
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.tools.tools_transform_service.dify_config") as mock_dify_config,
):
# Setup default mock returns
mock_dify_config.CONSOLE_API_URL = "https://console.example.com"
yield {
"dify_config": mock_dify_config,
}
def _create_test_tool_provider(
self, db_session_with_containers, mock_external_service_dependencies, provider_type="api"
):
"""
Helper method to create a test tool provider for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
mock_external_service_dependencies: Mock dependencies
provider_type: Type of provider to create
Returns:
Tool provider instance
"""
fake = Faker()
if provider_type == "api":
provider = ApiToolProvider(
name=fake.company(),
description=fake.text(max_nb_chars=100),
icon='{"background": "#FF6B6B", "content": "🔧"}',
icon_dark='{"background": "#252525", "content": "🔧"}',
tenant_id="test_tenant_id",
user_id="test_user_id",
credentials={"auth_type": "api_key_header", "api_key": "test_key"},
provider_type="api",
)
elif provider_type == "builtin":
provider = BuiltinToolProvider(
name=fake.company(),
description=fake.text(max_nb_chars=100),
icon="🔧",
icon_dark="🔧",
tenant_id="test_tenant_id",
provider="test_provider",
credential_type="api_key",
credentials={"api_key": "test_key"},
)
elif provider_type == "workflow":
provider = WorkflowToolProvider(
name=fake.company(),
description=fake.text(max_nb_chars=100),
icon='{"background": "#FF6B6B", "content": "🔧"}',
icon_dark='{"background": "#252525", "content": "🔧"}',
tenant_id="test_tenant_id",
user_id="test_user_id",
workflow_id="test_workflow_id",
)
elif provider_type == "mcp":
provider = MCPToolProvider(
name=fake.company(),
description=fake.text(max_nb_chars=100),
provider_icon='{"background": "#FF6B6B", "content": "🔧"}',
tenant_id="test_tenant_id",
user_id="test_user_id",
server_url="https://mcp.example.com",
server_identifier="test_server",
tools='[{"name": "test_tool", "description": "Test tool"}]',
authed=True,
)
else:
raise ValueError(f"Unknown provider type: {provider_type}")
from extensions.ext_database import db
db.session.add(provider)
db.session.commit()
return provider
def test_get_plugin_icon_url_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful plugin icon URL generation.
This test verifies:
- Proper URL construction for plugin icons
- Correct tenant_id and filename handling
- URL format compliance
"""
# Arrange: Setup test data
fake = Faker()
tenant_id = fake.uuid4()
filename = "test_icon.png"
# Act: Execute the method under test
result = ToolTransformService.get_plugin_icon_url(tenant_id, filename)
# Assert: Verify the expected outcomes
assert result is not None
assert isinstance(result, str)
assert "console/api/workspaces/current/plugin/icon" in result
assert tenant_id in result
assert filename in result
assert result.startswith("https://console.example.com")
# Verify URL structure
expected_url = f"https://console.example.com/console/api/workspaces/current/plugin/icon?tenant_id={tenant_id}&filename={filename}"
assert result == expected_url
def test_get_plugin_icon_url_with_empty_console_url(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test plugin icon URL generation when CONSOLE_API_URL is empty.
This test verifies:
- Fallback to relative URL when CONSOLE_API_URL is None
- Proper URL construction with relative path
"""
# Arrange: Setup mock with empty console URL
mock_external_service_dependencies["dify_config"].CONSOLE_API_URL = None
fake = Faker()
tenant_id = fake.uuid4()
filename = "test_icon.png"
# Act: Execute the method under test
result = ToolTransformService.get_plugin_icon_url(tenant_id, filename)
# Assert: Verify the expected outcomes
assert result is not None
assert isinstance(result, str)
assert result.startswith("/console/api/workspaces/current/plugin/icon")
assert tenant_id in result
assert filename in result
# Verify URL structure
expected_url = f"/console/api/workspaces/current/plugin/icon?tenant_id={tenant_id}&filename={filename}"
assert result == expected_url
def test_get_tool_provider_icon_url_builtin_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful tool provider icon URL generation for builtin providers.
This test verifies:
- Proper URL construction for builtin tool providers
- Correct provider type handling
- URL format compliance
"""
# Arrange: Setup test data
fake = Faker()
provider_type = ToolProviderType.BUILT_IN.value
provider_name = fake.company()
icon = "🔧"
# Act: Execute the method under test
result = ToolTransformService.get_tool_provider_icon_url(provider_type, provider_name, icon)
# Assert: Verify the expected outcomes
assert result is not None
assert isinstance(result, str)
assert "console/api/workspaces/current/tool-provider/builtin" in result
# Note: provider_name may contain spaces that get URL encoded
assert provider_name.replace(" ", "%20") in result or provider_name in result
assert result.endswith("/icon")
assert result.startswith("https://console.example.com")
# Verify URL structure (accounting for URL encoding)
# The actual result will have URL-encoded spaces (%20), so we need to compare accordingly
expected_url = (
f"https://console.example.com/console/api/workspaces/current/tool-provider/builtin/{provider_name}/icon"
)
# Convert expected URL to match the actual URL encoding
expected_encoded = expected_url.replace(" ", "%20")
assert result == expected_encoded
def test_get_tool_provider_icon_url_api_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful tool provider icon URL generation for API providers.
This test verifies:
- Proper icon handling for API tool providers
- JSON string parsing for icon data
- Fallback icon when parsing fails
"""
# Arrange: Setup test data
fake = Faker()
provider_type = ToolProviderType.API.value
provider_name = fake.company()
icon = '{"background": "#FF6B6B", "content": "🔧"}'
# Act: Execute the method under test
result = ToolTransformService.get_tool_provider_icon_url(provider_type, provider_name, icon)
# Assert: Verify the expected outcomes
assert result is not None
assert isinstance(result, dict)
assert result["background"] == "#FF6B6B"
assert result["content"] == "🔧"
def test_get_tool_provider_icon_url_api_invalid_json(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test tool provider icon URL generation for API providers with invalid JSON.
This test verifies:
- Proper fallback when JSON parsing fails
- Default icon structure when exception occurs
"""
# Arrange: Setup test data with invalid JSON
fake = Faker()
provider_type = ToolProviderType.API.value
provider_name = fake.company()
icon = '{"invalid": json}'
# Act: Execute the method under test
result = ToolTransformService.get_tool_provider_icon_url(provider_type, provider_name, icon)
# Assert: Verify the expected outcomes
assert result is not None
assert isinstance(result, dict)
assert result["background"] == "#252525"
# Note: emoji characters may be represented as Unicode escape sequences
assert result["content"] == "😁" or result["content"] == "\ud83d\ude01"
def test_get_tool_provider_icon_url_workflow_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful tool provider icon URL generation for workflow providers.
This test verifies:
- Proper icon handling for workflow tool providers
- Direct icon return for workflow type
"""
# Arrange: Setup test data
fake = Faker()
provider_type = ToolProviderType.WORKFLOW.value
provider_name = fake.company()
icon = {"background": "#FF6B6B", "content": "🔧"}
# Act: Execute the method under test
result = ToolTransformService.get_tool_provider_icon_url(provider_type, provider_name, icon)
# Assert: Verify the expected outcomes
assert result is not None
assert isinstance(result, dict)
assert result["background"] == "#FF6B6B"
assert result["content"] == "🔧"
def test_get_tool_provider_icon_url_mcp_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful tool provider icon URL generation for MCP providers.
This test verifies:
- Direct icon return for MCP type
- No URL transformation for MCP providers
"""
# Arrange: Setup test data
fake = Faker()
provider_type = ToolProviderType.MCP.value
provider_name = fake.company()
icon = {"background": "#FF6B6B", "content": "🔧"}
# Act: Execute the method under test
result = ToolTransformService.get_tool_provider_icon_url(provider_type, provider_name, icon)
# Assert: Verify the expected outcomes
assert result is not None
assert isinstance(result, dict)
assert result["background"] == "#FF6B6B"
assert result["content"] == "🔧"
def test_get_tool_provider_icon_url_unknown_type(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test tool provider icon URL generation for unknown provider types.
This test verifies:
- Empty string return for unknown provider types
- Proper handling of unsupported types
"""
# Arrange: Setup test data with unknown type
fake = Faker()
provider_type = "unknown_type"
provider_name = fake.company()
icon = "🔧"
# Act: Execute the method under test
result = ToolTransformService.get_tool_provider_icon_url(provider_type, provider_name, icon)
# Assert: Verify the expected outcomes
assert result == ""
def test_repack_provider_dict_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful provider repacking with dictionary input.
This test verifies:
- Proper icon URL generation for dictionary providers
- Correct provider type handling
- Icon transformation for different provider types
"""
# Arrange: Setup test data
fake = Faker()
tenant_id = fake.uuid4()
provider = {"type": ToolProviderType.BUILT_IN.value, "name": fake.company(), "icon": "🔧"}
# Act: Execute the method under test
ToolTransformService.repack_provider(tenant_id, provider)
# Assert: Verify the expected outcomes
assert "icon" in provider
assert isinstance(provider["icon"], str)
assert "console/api/workspaces/current/tool-provider/builtin" in provider["icon"]
# Note: provider name may contain spaces that get URL encoded
assert provider["name"].replace(" ", "%20") in provider["icon"] or provider["name"] in provider["icon"]
def test_repack_provider_entity_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful provider repacking with ToolProviderApiEntity input.
This test verifies:
- Proper icon URL generation for entity providers
- Plugin icon handling when plugin_id is present
- Regular icon handling when plugin_id is not present
"""
# Arrange: Setup test data
fake = Faker()
tenant_id = fake.uuid4()
# Create provider entity with plugin_id
provider = ToolProviderApiEntity(
id=fake.uuid4(),
author=fake.name(),
name=fake.company(),
description=I18nObject(en_US=fake.text(max_nb_chars=100)),
icon="test_icon.png",
icon_dark="test_icon_dark.png",
label=I18nObject(en_US=fake.company()),
type=ToolProviderType.API,
masked_credentials={},
is_team_authorization=True,
plugin_id="test_plugin_id",
tools=[],
labels=[],
)
# Act: Execute the method under test
ToolTransformService.repack_provider(tenant_id, provider)
# Assert: Verify the expected outcomes
assert provider.icon is not None
assert isinstance(provider.icon, str)
assert "console/api/workspaces/current/plugin/icon" in provider.icon
assert tenant_id in provider.icon
assert "test_icon.png" in provider.icon
# Verify dark icon handling
assert provider.icon_dark is not None
assert isinstance(provider.icon_dark, str)
assert "console/api/workspaces/current/plugin/icon" in provider.icon_dark
assert tenant_id in provider.icon_dark
assert "test_icon_dark.png" in provider.icon_dark
def test_repack_provider_entity_no_plugin_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful provider repacking with ToolProviderApiEntity input without plugin_id.
This test verifies:
- Proper icon URL generation for non-plugin providers
- Regular tool provider icon handling
- Dark icon handling when present
"""
# Arrange: Setup test data
fake = Faker()
tenant_id = fake.uuid4()
# Create provider entity without plugin_id
provider = ToolProviderApiEntity(
id=fake.uuid4(),
author=fake.name(),
name=fake.company(),
description=I18nObject(en_US=fake.text(max_nb_chars=100)),
icon='{"background": "#FF6B6B", "content": "🔧"}',
icon_dark='{"background": "#252525", "content": "🔧"}',
label=I18nObject(en_US=fake.company()),
type=ToolProviderType.API,
masked_credentials={},
is_team_authorization=True,
plugin_id=None,
tools=[],
labels=[],
)
# Act: Execute the method under test
ToolTransformService.repack_provider(tenant_id, provider)
# Assert: Verify the expected outcomes
assert provider.icon is not None
assert isinstance(provider.icon, dict)
assert provider.icon["background"] == "#FF6B6B"
assert provider.icon["content"] == "🔧"
# Verify dark icon handling
assert provider.icon_dark is not None
assert isinstance(provider.icon_dark, dict)
assert provider.icon_dark["background"] == "#252525"
assert provider.icon_dark["content"] == "🔧"
def test_repack_provider_entity_no_dark_icon(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test provider repacking with ToolProviderApiEntity input without dark icon.
This test verifies:
- Proper handling when icon_dark is None or empty
- No errors when dark icon is not present
"""
# Arrange: Setup test data
fake = Faker()
tenant_id = fake.uuid4()
# Create provider entity without dark icon
provider = ToolProviderApiEntity(
id=fake.uuid4(),
author=fake.name(),
name=fake.company(),
description=I18nObject(en_US=fake.text(max_nb_chars=100)),
icon='{"background": "#FF6B6B", "content": "🔧"}',
icon_dark="",
label=I18nObject(en_US=fake.company()),
type=ToolProviderType.API,
masked_credentials={},
is_team_authorization=True,
plugin_id=None,
tools=[],
labels=[],
)
# Act: Execute the method under test
ToolTransformService.repack_provider(tenant_id, provider)
# Assert: Verify the expected outcomes
assert provider.icon is not None
assert isinstance(provider.icon, dict)
assert provider.icon["background"] == "#FF6B6B"
assert provider.icon["content"] == "🔧"
# Verify dark icon remains empty string
assert provider.icon_dark == ""
def test_builtin_provider_to_user_provider_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful conversion of builtin provider to user provider.
This test verifies:
- Proper entity creation with all required fields
- Credentials schema handling
- Team authorization setup
- Plugin ID handling
"""
# Arrange: Setup test data
fake = Faker()
# Create mock provider controller
mock_controller = Mock()
mock_controller.entity.identity.name = fake.company()
mock_controller.entity.identity.author = fake.name()
mock_controller.entity.identity.description = I18nObject(en_US=fake.text(max_nb_chars=100))
mock_controller.entity.identity.icon = "🔧"
mock_controller.entity.identity.icon_dark = "🔧"
mock_controller.entity.identity.label = I18nObject(en_US=fake.company())
mock_controller.plugin_id = None
mock_controller.plugin_unique_identifier = None
mock_controller.tool_labels = ["label1", "label2"]
mock_controller.need_credentials = True
# Mock credentials schema
mock_credential = Mock()
mock_credential.to_basic_provider_config.return_value.name = "api_key"
mock_controller.get_credentials_schema_by_type.return_value = [mock_credential]
# Create mock database provider
mock_db_provider = Mock()
mock_db_provider.credential_type = "api-key"
mock_db_provider.tenant_id = fake.uuid4()
mock_db_provider.credentials = {"api_key": "encrypted_key"}
# Mock encryption
with patch("services.tools.tools_transform_service.create_provider_encrypter") as mock_encrypter:
mock_encrypter_instance = Mock()
mock_encrypter_instance.decrypt.return_value = {"api_key": "decrypted_key"}
mock_encrypter_instance.mask_tool_credentials.return_value = {"api_key": ""}
mock_encrypter.return_value = (mock_encrypter_instance, None)
# Act: Execute the method under test
result = ToolTransformService.builtin_provider_to_user_provider(
mock_controller, mock_db_provider, decrypt_credentials=True
)
# Assert: Verify the expected outcomes
assert result is not None
assert result.id == mock_controller.entity.identity.name
assert result.author == mock_controller.entity.identity.author
assert result.name == mock_controller.entity.identity.name
assert result.description == mock_controller.entity.identity.description
assert result.icon == mock_controller.entity.identity.icon
assert result.icon_dark == mock_controller.entity.identity.icon_dark
assert result.label == mock_controller.entity.identity.label
assert result.type == ToolProviderType.BUILT_IN
assert result.is_team_authorization is True
assert result.plugin_id is None
assert result.tools == []
assert result.labels == ["label1", "label2"]
assert result.masked_credentials == {"api_key": ""}
assert result.original_credentials == {"api_key": "decrypted_key"}
def test_builtin_provider_to_user_provider_plugin_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful conversion of builtin provider to user provider with plugin.
This test verifies:
- Plugin ID and unique identifier handling
- Proper entity creation for plugin providers
"""
# Arrange: Setup test data
fake = Faker()
# Create mock provider controller with plugin
mock_controller = Mock()
mock_controller.entity.identity.name = fake.company()
mock_controller.entity.identity.author = fake.name()
mock_controller.entity.identity.description = I18nObject(en_US=fake.text(max_nb_chars=100))
mock_controller.entity.identity.icon = "🔧"
mock_controller.entity.identity.icon_dark = "🔧"
mock_controller.entity.identity.label = I18nObject(en_US=fake.company())
mock_controller.plugin_id = "test_plugin_id"
mock_controller.plugin_unique_identifier = "test_unique_id"
mock_controller.tool_labels = ["label1"]
mock_controller.need_credentials = False
# Mock credentials schema
mock_credential = Mock()
mock_credential.to_basic_provider_config.return_value.name = "api_key"
mock_controller.get_credentials_schema_by_type.return_value = [mock_credential]
# Act: Execute the method under test
result = ToolTransformService.builtin_provider_to_user_provider(
mock_controller, None, decrypt_credentials=False
)
# Assert: Verify the expected outcomes
assert result is not None
# Note: The method checks isinstance(provider_controller, PluginToolProviderController)
# Since we're using a Mock, this check will fail, so plugin_id will remain None
# In a real test with actual PluginToolProviderController, this would work
assert result.is_team_authorization is True
assert result.allow_delete is False
def test_builtin_provider_to_user_provider_no_credentials(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test conversion of builtin provider to user provider without credentials.
This test verifies:
- Proper handling when no credentials are needed
- Team authorization setup for no-credentials providers
"""
# Arrange: Setup test data
fake = Faker()
# Create mock provider controller
mock_controller = Mock()
mock_controller.entity.identity.name = fake.company()
mock_controller.entity.identity.author = fake.name()
mock_controller.entity.identity.description = I18nObject(en_US=fake.text(max_nb_chars=100))
mock_controller.entity.identity.icon = "🔧"
mock_controller.entity.identity.icon_dark = "🔧"
mock_controller.entity.identity.label = I18nObject(en_US=fake.company())
mock_controller.plugin_id = None
mock_controller.plugin_unique_identifier = None
mock_controller.tool_labels = []
mock_controller.need_credentials = False
# Mock credentials schema
mock_credential = Mock()
mock_credential.to_basic_provider_config.return_value.name = "api_key"
mock_controller.get_credentials_schema_by_type.return_value = [mock_credential]
# Act: Execute the method under test
result = ToolTransformService.builtin_provider_to_user_provider(
mock_controller, None, decrypt_credentials=False
)
# Assert: Verify the expected outcomes
assert result is not None
assert result.is_team_authorization is True
assert result.allow_delete is False
assert result.masked_credentials == {"api_key": ""}
def test_api_provider_to_controller_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful conversion of API provider to controller.
This test verifies:
- Proper controller creation from database provider
- Auth type handling for different credential types
- Backward compatibility for auth types
"""
# Arrange: Setup test data
fake = Faker()
# Create API tool provider with api_key_header auth
provider = ApiToolProvider(
name=fake.company(),
description=fake.text(max_nb_chars=100),
icon='{"background": "#FF6B6B", "content": "🔧"}',
tenant_id=fake.uuid4(),
user_id=fake.uuid4(),
credentials_str='{"auth_type": "api_key_header", "api_key": "test_key"}',
schema="{}",
schema_type_str="openapi",
tools_str="[]",
)
from extensions.ext_database import db
db.session.add(provider)
db.session.commit()
# Act: Execute the method under test
result = ToolTransformService.api_provider_to_controller(provider)
# Assert: Verify the expected outcomes
assert result is not None
assert hasattr(result, "from_db")
# Additional assertions would depend on the actual controller implementation
def test_api_provider_to_controller_api_key_query(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test conversion of API provider to controller with api_key_query auth type.
This test verifies:
- Proper auth type handling for query parameter authentication
"""
# Arrange: Setup test data
fake = Faker()
# Create API tool provider with api_key_query auth
provider = ApiToolProvider(
name=fake.company(),
description=fake.text(max_nb_chars=100),
icon='{"background": "#FF6B6B", "content": "🔧"}',
tenant_id=fake.uuid4(),
user_id=fake.uuid4(),
credentials_str='{"auth_type": "api_key_query", "api_key": "test_key"}',
schema="{}",
schema_type_str="openapi",
tools_str="[]",
)
from extensions.ext_database import db
db.session.add(provider)
db.session.commit()
# Act: Execute the method under test
result = ToolTransformService.api_provider_to_controller(provider)
# Assert: Verify the expected outcomes
assert result is not None
assert hasattr(result, "from_db")
def test_api_provider_to_controller_backward_compatibility(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test conversion of API provider to controller with backward compatibility auth types.
This test verifies:
- Proper handling of legacy auth type values
- Backward compatibility for api_key and api_key_header
"""
# Arrange: Setup test data
fake = Faker()
# Create API tool provider with legacy auth type
provider = ApiToolProvider(
name=fake.company(),
description=fake.text(max_nb_chars=100),
icon='{"background": "#FF6B6B", "content": "🔧"}',
tenant_id=fake.uuid4(),
user_id=fake.uuid4(),
credentials_str='{"auth_type": "api_key", "api_key": "test_key"}',
schema="{}",
schema_type_str="openapi",
tools_str="[]",
)
from extensions.ext_database import db
db.session.add(provider)
db.session.commit()
# Act: Execute the method under test
result = ToolTransformService.api_provider_to_controller(provider)
# Assert: Verify the expected outcomes
assert result is not None
assert hasattr(result, "from_db")
def test_workflow_provider_to_controller_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful conversion of workflow provider to controller.
This test verifies:
- Proper controller creation from workflow provider
- Workflow-specific controller handling
"""
# Arrange: Setup test data
fake = Faker()
# Create workflow tool provider
provider = WorkflowToolProvider(
name=fake.company(),
description=fake.text(max_nb_chars=100),
icon='{"background": "#FF6B6B", "content": "🔧"}',
tenant_id=fake.uuid4(),
user_id=fake.uuid4(),
app_id=fake.uuid4(),
label="Test Workflow",
version="1.0.0",
parameter_configuration="[]",
)
from extensions.ext_database import db
db.session.add(provider)
db.session.commit()
# Mock the WorkflowToolProviderController.from_db method to avoid app dependency
with patch("services.tools.tools_transform_service.WorkflowToolProviderController.from_db") as mock_from_db:
mock_controller = Mock()
mock_from_db.return_value = mock_controller
# Act: Execute the method under test
result = ToolTransformService.workflow_provider_to_controller(provider)
# Assert: Verify the expected outcomes
assert result is not None
assert result == mock_controller
mock_from_db.assert_called_once_with(provider)

716
api/tests/test_containers_integration_tests/services/tools/test_workflow_tools_manage_service.py Normal file
View File

@ -0,0 +1,716 @@
import json
from unittest.mock import patch
import pytest
from faker import Faker
from models.tools import WorkflowToolProvider
from models.workflow import Workflow as WorkflowModel
from services.account_service import AccountService, TenantService
from services.app_service import AppService
from services.tools.workflow_tools_manage_service import WorkflowToolManageService
class TestWorkflowToolManageService:
"""Integration tests for WorkflowToolManageService using testcontainers."""
@pytest.fixture
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
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.tools.workflow_tools_manage_service.WorkflowToolProviderController"
) as mock_workflow_tool_provider_controller,
patch("services.tools.workflow_tools_manage_service.ToolLabelManager") as mock_tool_label_manager,
patch("services.tools.workflow_tools_manage_service.ToolTransformService") as mock_tool_transform_service,
):
# Setup default mock returns for app service
mock_feature_service.get_system_features.return_value.webapp_auth.enabled = False
mock_enterprise_service.WebAppAuth.update_app_access_mode.return_value = None
mock_enterprise_service.WebAppAuth.cleanup_webapp.return_value = None
# Setup default mock returns for account service
mock_account_feature_service.get_system_features.return_value.is_allow_register = True
# Mock ModelManager for model configuration
mock_model_instance = mock_model_manager.return_value
mock_model_instance.get_default_model_instance.return_value = None
mock_model_instance.get_default_provider_model_name.return_value = ("openai", "gpt-3.5-turbo")
# Mock WorkflowToolProviderController
mock_workflow_tool_provider_controller.from_db.return_value = None
# Mock ToolLabelManager
mock_tool_label_manager.update_tool_labels.return_value = None
# Mock ToolTransformService
mock_tool_transform_service.workflow_provider_to_controller.return_value = None
yield {
"feature_service": mock_feature_service,
"enterprise_service": mock_enterprise_service,
"model_manager": mock_model_manager,
"account_feature_service": mock_account_feature_service,
"workflow_tool_provider_controller": mock_workflow_tool_provider_controller,
"tool_label_manager": mock_tool_label_manager,
"tool_transform_service": mock_tool_transform_service,
}
def _create_test_app_and_account(self, db_session_with_containers, mock_external_service_dependencies):
"""
Helper method to create a test app and account for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
mock_external_service_dependencies: Mock dependencies
Returns:
tuple: (app, account, workflow) - Created app, account and workflow instances
"""
fake = Faker()
# Setup mocks for account creation
mock_external_service_dependencies[
"account_feature_service"
].get_system_features.return_value.is_allow_register = True
# Create account and tenant
account = AccountService.create_account(
email=fake.email(),
name=fake.name(),
interface_language="en-US",
password=fake.password(length=12),
)
TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
tenant = account.current_tenant
# Create app with realistic data
app_args = {
"name": fake.company(),
"description": fake.text(max_nb_chars=100),
"mode": "workflow",
"icon_type": "emoji",
"icon": "🤖",
"icon_background": "#FF6B6B",
"api_rph": 100,
"api_rpm": 10,
}
app_service = AppService()
app = app_service.create_app(tenant.id, app_args, account)
# Create workflow for the app
workflow = WorkflowModel(
tenant_id=tenant.id,
app_id=app.id,
type="workflow",
version="1.0.0",
graph=json.dumps({}),
features=json.dumps({}),
created_by=account.id,
environment_variables=[],
conversation_variables=[],
)
from extensions.ext_database import db
db.session.add(workflow)
db.session.commit()
# Update app to reference the workflow
app.workflow_id = workflow.id
db.session.commit()
return app, account, workflow
def _create_test_workflow_tool_parameters(self):
"""Helper method to create valid workflow tool parameters."""
return [
{
"name": "input_text",
"description": "Input text for processing",
"form": "form",
"type": "string",
"required": True,
},
{
"name": "output_format",
"description": "Output format specification",
"form": "form",
"type": "select",
"required": False,
},
]
def test_create_workflow_tool_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful workflow tool creation with valid parameters.
This test verifies:
- Proper workflow tool creation with all required fields
- Correct database state after creation
- Proper relationship establishment
- External service integration
- Return value correctness
"""
fake = Faker()
# Create test data
app, account, workflow = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Setup workflow tool creation parameters
tool_name = fake.word()
tool_label = fake.word()
tool_icon = {"type": "emoji", "emoji": "🔧"}
tool_description = fake.text(max_nb_chars=200)
tool_parameters = self._create_test_workflow_tool_parameters()
tool_privacy_policy = fake.text(max_nb_chars=100)
tool_labels = ["automation", "workflow"]
# Execute the method under test
result = WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id,
name=tool_name,
label=tool_label,
icon=tool_icon,
description=tool_description,
parameters=tool_parameters,
privacy_policy=tool_privacy_policy,
labels=tool_labels,
)
# Verify the result
assert result == {"result": "success"}
# Verify database state
from extensions.ext_database import db
# Check if workflow tool provider was created
created_tool_provider = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
WorkflowToolProvider.app_id == app.id,
)
.first()
)
assert created_tool_provider is not None
assert created_tool_provider.name == tool_name
assert created_tool_provider.label == tool_label
assert created_tool_provider.icon == json.dumps(tool_icon)
assert created_tool_provider.description == tool_description
assert created_tool_provider.parameter_configuration == json.dumps(tool_parameters)
assert created_tool_provider.privacy_policy == tool_privacy_policy
assert created_tool_provider.version == workflow.version
assert created_tool_provider.user_id == account.id
assert created_tool_provider.tenant_id == account.current_tenant.id
assert created_tool_provider.app_id == app.id
# Verify external service calls
mock_external_service_dependencies["workflow_tool_provider_controller"].from_db.assert_called_once()
mock_external_service_dependencies["tool_label_manager"].update_tool_labels.assert_called_once()
mock_external_service_dependencies[
"tool_transform_service"
].workflow_provider_to_controller.assert_called_once()
def test_create_workflow_tool_duplicate_name_error(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test workflow tool creation fails when name already exists.
This test verifies:
- Proper error handling for duplicate tool names
- Database constraint enforcement
- Correct error message
"""
fake = Faker()
# Create test data
app, account, workflow = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Create first workflow tool
first_tool_name = fake.word()
first_tool_parameters = self._create_test_workflow_tool_parameters()
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id,
name=first_tool_name,
label=fake.word(),
icon={"type": "emoji", "emoji": "🔧"},
description=fake.text(max_nb_chars=200),
parameters=first_tool_parameters,
)
# Attempt to create second workflow tool with same name
second_tool_parameters = self._create_test_workflow_tool_parameters()
with pytest.raises(ValueError) as exc_info:
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id,
name=first_tool_name, # Same name
label=fake.word(),
icon={"type": "emoji", "emoji": "⚙️"},
description=fake.text(max_nb_chars=200),
parameters=second_tool_parameters,
)
# Verify error message
assert f"Tool with name {first_tool_name} or app_id {app.id} already exists" in str(exc_info.value)
# Verify only one tool was created
from extensions.ext_database import db
tool_count = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
)
.count()
)
assert tool_count == 1
def test_create_workflow_tool_invalid_app_error(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test workflow tool creation fails when app does not exist.
This test verifies:
- Proper error handling for non-existent apps
- Correct error message
- No database changes when app is invalid
"""
fake = Faker()
# Create test data
app, account, workflow = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Generate non-existent app ID
non_existent_app_id = fake.uuid4()
# Attempt to create workflow tool with non-existent app
tool_parameters = self._create_test_workflow_tool_parameters()
with pytest.raises(ValueError) as exc_info:
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=non_existent_app_id, # Non-existent app ID
name=fake.word(),
label=fake.word(),
icon={"type": "emoji", "emoji": "🔧"},
description=fake.text(max_nb_chars=200),
parameters=tool_parameters,
)
# Verify error message
assert f"App {non_existent_app_id} not found" in str(exc_info.value)
# Verify no workflow tool was created
from extensions.ext_database import db
tool_count = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
)
.count()
)
assert tool_count == 0
def test_create_workflow_tool_invalid_parameters_error(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test workflow tool creation fails when parameters are invalid.
This test verifies:
- Proper error handling for invalid parameter configurations
- Parameter validation enforcement
- Correct error message
"""
fake = Faker()
# Create test data
app, account, workflow = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Setup invalid workflow tool parameters (missing required fields)
invalid_parameters = [
{
"name": "input_text",
# Missing description and form fields
"type": "string",
"required": True,
}
]
# Attempt to create workflow tool with invalid parameters
with pytest.raises(ValueError) as exc_info:
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id,
name=fake.word(),
label=fake.word(),
icon={"type": "emoji", "emoji": "🔧"},
description=fake.text(max_nb_chars=200),
parameters=invalid_parameters,
)
# Verify error message contains validation error
assert "validation error" in str(exc_info.value).lower()
# Verify no workflow tool was created
from extensions.ext_database import db
tool_count = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
)
.count()
)
assert tool_count == 0
def test_create_workflow_tool_duplicate_app_id_error(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test workflow tool creation fails when app_id already exists.
This test verifies:
- Proper error handling for duplicate app_id
- Database constraint enforcement for app_id uniqueness
- Correct error message
"""
fake = Faker()
# Create test data
app, account, workflow = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Create first workflow tool
first_tool_name = fake.word()
first_tool_parameters = self._create_test_workflow_tool_parameters()
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id,
name=first_tool_name,
label=fake.word(),
icon={"type": "emoji", "emoji": "🔧"},
description=fake.text(max_nb_chars=200),
parameters=first_tool_parameters,
)
# Attempt to create second workflow tool with same app_id but different name
second_tool_name = fake.word()
second_tool_parameters = self._create_test_workflow_tool_parameters()
with pytest.raises(ValueError) as exc_info:
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id, # Same app_id
name=second_tool_name, # Different name
label=fake.word(),
icon={"type": "emoji", "emoji": "⚙️"},
description=fake.text(max_nb_chars=200),
parameters=second_tool_parameters,
)
# Verify error message
assert f"Tool with name {second_tool_name} or app_id {app.id} already exists" in str(exc_info.value)
# Verify only one tool was created
from extensions.ext_database import db
tool_count = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
)
.count()
)
assert tool_count == 1
def test_create_workflow_tool_workflow_not_found_error(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test workflow tool creation fails when app has no workflow.
This test verifies:
- Proper error handling for apps without workflows
- Correct error message
- No database changes when workflow is missing
"""
fake = Faker()
# Create test data but without workflow
app, account, _ = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Remove workflow reference from app
from extensions.ext_database import db
app.workflow_id = None
db.session.commit()
# Attempt to create workflow tool for app without workflow
tool_parameters = self._create_test_workflow_tool_parameters()
with pytest.raises(ValueError) as exc_info:
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id,
name=fake.word(),
label=fake.word(),
icon={"type": "emoji", "emoji": "🔧"},
description=fake.text(max_nb_chars=200),
parameters=tool_parameters,
)
# Verify error message
assert f"Workflow not found for app {app.id}" in str(exc_info.value)
# Verify no workflow tool was created
tool_count = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
)
.count()
)
assert tool_count == 0
def test_update_workflow_tool_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful workflow tool update with valid parameters.
This test verifies:
- Proper workflow tool update with all required fields
- Correct database state after update
- Proper relationship maintenance
- External service integration
- Return value correctness
"""
fake = Faker()
# Create test data
app, account, workflow = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Create initial workflow tool
initial_tool_name = fake.word()
initial_tool_parameters = self._create_test_workflow_tool_parameters()
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id,
name=initial_tool_name,
label=fake.word(),
icon={"type": "emoji", "emoji": "🔧"},
description=fake.text(max_nb_chars=200),
parameters=initial_tool_parameters,
)
# Get the created tool
from extensions.ext_database import db
created_tool = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
WorkflowToolProvider.app_id == app.id,
)
.first()
)
# Setup update parameters
updated_tool_name = fake.word()
updated_tool_label = fake.word()
updated_tool_icon = {"type": "emoji", "emoji": "⚙️"}
updated_tool_description = fake.text(max_nb_chars=200)
updated_tool_parameters = self._create_test_workflow_tool_parameters()
updated_tool_privacy_policy = fake.text(max_nb_chars=100)
updated_tool_labels = ["automation", "updated"]
# Execute the update method
result = WorkflowToolManageService.update_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_tool_id=created_tool.id,
name=updated_tool_name,
label=updated_tool_label,
icon=updated_tool_icon,
description=updated_tool_description,
parameters=updated_tool_parameters,
privacy_policy=updated_tool_privacy_policy,
labels=updated_tool_labels,
)
# Verify the result
assert result == {"result": "success"}
# Verify database state was updated
db.session.refresh(created_tool)
assert created_tool.name == updated_tool_name
assert created_tool.label == updated_tool_label
assert created_tool.icon == json.dumps(updated_tool_icon)
assert created_tool.description == updated_tool_description
assert created_tool.parameter_configuration == json.dumps(updated_tool_parameters)
assert created_tool.privacy_policy == updated_tool_privacy_policy
assert created_tool.version == workflow.version
assert created_tool.updated_at is not None
# Verify external service calls
mock_external_service_dependencies["workflow_tool_provider_controller"].from_db.assert_called()
mock_external_service_dependencies["tool_label_manager"].update_tool_labels.assert_called()
mock_external_service_dependencies["tool_transform_service"].workflow_provider_to_controller.assert_called()
def test_update_workflow_tool_not_found_error(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test workflow tool update fails when tool does not exist.
This test verifies:
- Proper error handling for non-existent tools
- Correct error message
- No database changes when tool is invalid
"""
fake = Faker()
# Create test data
app, account, workflow = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Generate non-existent tool ID
non_existent_tool_id = fake.uuid4()
# Attempt to update non-existent workflow tool
tool_parameters = self._create_test_workflow_tool_parameters()
with pytest.raises(ValueError) as exc_info:
WorkflowToolManageService.update_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_tool_id=non_existent_tool_id, # Non-existent tool ID
name=fake.word(),
label=fake.word(),
icon={"type": "emoji", "emoji": "🔧"},
description=fake.text(max_nb_chars=200),
parameters=tool_parameters,
)
# Verify error message
assert f"Tool {non_existent_tool_id} not found" in str(exc_info.value)
# Verify no workflow tool was created
from extensions.ext_database import db
tool_count = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
)
.count()
)
assert tool_count == 0
def test_update_workflow_tool_same_name_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test workflow tool update succeeds when keeping the same name.
This test verifies:
- Proper handling when updating tool with same name
- Database state maintenance
- Update timestamp is set
"""
fake = Faker()
# Create test data
app, account, workflow = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies
)
# Create first workflow tool
first_tool_name = fake.word()
first_tool_parameters = self._create_test_workflow_tool_parameters()
WorkflowToolManageService.create_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_app_id=app.id,
name=first_tool_name,
label=fake.word(),
icon={"type": "emoji", "emoji": "🔧"},
description=fake.text(max_nb_chars=200),
parameters=first_tool_parameters,
)
# Get the created tool
from extensions.ext_database import db
created_tool = (
db.session.query(WorkflowToolProvider)
.where(
WorkflowToolProvider.tenant_id == account.current_tenant.id,
WorkflowToolProvider.app_id == app.id,
)
.first()
)
# Attempt to update tool with same name (should not fail)
result = WorkflowToolManageService.update_workflow_tool(
user_id=account.id,
tenant_id=account.current_tenant.id,
workflow_tool_id=created_tool.id,
name=first_tool_name, # Same name
label=fake.word(),
icon={"type": "emoji", "emoji": "⚙️"},
description=fake.text(max_nb_chars=200),
parameters=first_tool_parameters,
)
# Verify update was successful
assert result == {"result": "success"}
# Verify tool still exists with the same name
db.session.refresh(created_tool)
assert created_tool.name == first_tool_name
assert created_tool.updated_at is not None

0
api/tests/test_containers_integration_tests/services/workflow/__init__.py Normal file
View File

554
api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py Normal file
View File

@ -0,0 +1,554 @@
import json
from unittest.mock import patch
import pytest
from faker import Faker
from core.app.app_config.entities import (
DatasetEntity,
DatasetRetrieveConfigEntity,
ExternalDataVariableEntity,
ModelConfigEntity,
PromptTemplateEntity,
VariableEntity,
VariableEntityType,
)
from core.model_runtime.entities.llm_entities import LLMMode
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from models.account import Account, Tenant
from models.api_based_extension import APIBasedExtension
from models.model import App, AppMode, AppModelConfig
from models.workflow import Workflow
from services.workflow.workflow_converter import WorkflowConverter
class TestWorkflowConverter:
"""Integration tests for WorkflowConverter using testcontainers."""
@pytest.fixture
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.workflow.workflow_converter.encrypter") as mock_encrypter,
patch("services.workflow.workflow_converter.SimplePromptTransform") as mock_prompt_transform,
patch("services.workflow.workflow_converter.AgentChatAppConfigManager") as mock_agent_chat_config_manager,
patch("services.workflow.workflow_converter.ChatAppConfigManager") as mock_chat_config_manager,
patch("services.workflow.workflow_converter.CompletionAppConfigManager") as mock_completion_config_manager,
):
# Setup default mock returns
mock_encrypter.decrypt_token.return_value = "decrypted_api_key"
mock_prompt_transform.return_value.get_prompt_template.return_value = {
"prompt_template": PromptTemplateParser(template="You are a helpful assistant {{text_input}}"),
"prompt_rules": {"human_prefix": "Human", "assistant_prefix": "Assistant"},
}
mock_agent_chat_config_manager.get_app_config.return_value = self._create_mock_app_config()
mock_chat_config_manager.get_app_config.return_value = self._create_mock_app_config()
mock_completion_config_manager.get_app_config.return_value = self._create_mock_app_config()
yield {
"encrypter": mock_encrypter,
"prompt_transform": mock_prompt_transform,
"agent_chat_config_manager": mock_agent_chat_config_manager,
"chat_config_manager": mock_chat_config_manager,
"completion_config_manager": mock_completion_config_manager,
}
def _create_mock_app_config(self):
"""Helper method to create a mock app config."""
mock_config = type("obj", (object,), {})()
mock_config.variables = [
VariableEntity(
variable="text_input",
label="Text Input",
type=VariableEntityType.TEXT_INPUT,
)
]
mock_config.model = ModelConfigEntity(
provider="openai",
model="gpt-4",
mode=LLMMode.CHAT.value,
parameters={},
stop=[],
)
mock_config.prompt_template = PromptTemplateEntity(
prompt_type=PromptTemplateEntity.PromptType.SIMPLE,
simple_prompt_template="You are a helpful assistant {{text_input}}",
)
mock_config.dataset = None
mock_config.external_data_variables = []
mock_config.additional_features = type("obj", (object,), {"file_upload": None})()
mock_config.app_model_config_dict = {}
return mock_config
def _create_test_account_and_tenant(self, db_session_with_containers, mock_external_service_dependencies):
"""
Helper method to create a test account and tenant for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
mock_external_service_dependencies: Mock dependencies
Returns:
tuple: (account, tenant) - Created account and tenant instances
"""
fake = Faker()
# Create account
account = Account(
email=fake.email(),
name=fake.name(),
interface_language="en-US",
status="active",
)
from extensions.ext_database import db
db.session.add(account)
db.session.commit()
# Create tenant for the account
tenant = Tenant(
name=fake.company(),
status="normal",
)
db.session.add(tenant)
db.session.commit()
# Create tenant-account join
from models.account import TenantAccountJoin, TenantAccountRole
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER.value,
current=True,
)
db.session.add(join)
db.session.commit()
# Set current tenant for account
account.current_tenant = tenant
return account, tenant
def _create_test_app(self, db_session_with_containers, mock_external_service_dependencies, tenant, account):
"""
Helper method to create a test app for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
mock_external_service_dependencies: Mock dependencies
tenant: Tenant instance
account: Account instance
Returns:
App: Created app instance
"""
fake = Faker()
# Create app
app = App(
tenant_id=tenant.id,
name=fake.company(),
mode=AppMode.CHAT.value,
icon_type="emoji",
icon="🤖",
icon_background="#FF6B6B",
enable_site=True,
enable_api=True,
api_rpm=100,
api_rph=10,
is_demo=False,
is_public=False,
created_by=account.id,
updated_by=account.id,
)
from extensions.ext_database import db
db.session.add(app)
db.session.commit()
# Create app model config
app_model_config = AppModelConfig(
app_id=app.id,
provider="openai",
model="gpt-4",
configs={},
created_by=account.id,
updated_by=account.id,
)
db.session.add(app_model_config)
db.session.commit()
# Link app model config to app
app.app_model_config_id = app_model_config.id
db.session.commit()
return app
def test_convert_to_workflow_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful conversion of app to workflow.
This test verifies:
- Proper app to workflow conversion
- Correct database state after conversion
- Proper relationship establishment
- Workflow creation with correct configuration
"""
# Arrange: Create test data
fake = Faker()
account, tenant = self._create_test_account_and_tenant(
db_session_with_containers, mock_external_service_dependencies
)
app = self._create_test_app(db_session_with_containers, mock_external_service_dependencies, tenant, account)
# Act: Execute the conversion
workflow_converter = WorkflowConverter()
new_app = workflow_converter.convert_to_workflow(
app_model=app,
account=account,
name="Test Workflow App",
icon_type="emoji",
icon="🚀",
icon_background="#4CAF50",
)
# Assert: Verify the expected outcomes
assert new_app is not None
assert new_app.name == "Test Workflow App"
assert new_app.mode == AppMode.ADVANCED_CHAT.value
assert new_app.icon_type == "emoji"
assert new_app.icon == "🚀"
assert new_app.icon_background == "#4CAF50"
assert new_app.tenant_id == app.tenant_id
assert new_app.created_by == account.id
# Verify database state
from extensions.ext_database import db
db.session.refresh(new_app)
assert new_app.id is not None
# Verify workflow was created
workflow = db.session.query(Workflow).where(Workflow.app_id == new_app.id).first()
assert workflow is not None
assert workflow.tenant_id == app.tenant_id
assert workflow.type == "chat"
def test_convert_to_workflow_without_app_model_config_error(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test error handling when app model config is missing.
This test verifies:
- Proper error handling for missing app model config
- Correct exception type and message
- Database state remains unchanged
"""
# Arrange: Create test data without app model config
fake = Faker()
account, tenant = self._create_test_account_and_tenant(
db_session_with_containers, mock_external_service_dependencies
)
app = App(
tenant_id=tenant.id,
name=fake.company(),
mode=AppMode.CHAT.value,
icon_type="emoji",
icon="🤖",
icon_background="#FF6B6B",
enable_site=True,
enable_api=True,
api_rpm=100,
api_rph=10,
is_demo=False,
is_public=False,
created_by=account.id,
updated_by=account.id,
)
from extensions.ext_database import db
db.session.add(app)
db.session.commit()
# Act & Assert: Verify proper error handling
workflow_converter = WorkflowConverter()
# Check initial state
initial_workflow_count = db.session.query(Workflow).count()
with pytest.raises(ValueError, match="App model config is required"):
workflow_converter.convert_to_workflow(
app_model=app,
account=account,
name="Test Workflow App",
icon_type="emoji",
icon="🚀",
icon_background="#4CAF50",
)
# Verify database state remains unchanged
# The workflow creation happens in convert_app_model_config_to_workflow
# which is called before the app_model_config check, so we need to clean up
db.session.rollback()
final_workflow_count = db.session.query(Workflow).count()
assert final_workflow_count == initial_workflow_count
def test_convert_app_model_config_to_workflow_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful conversion of app model config to workflow.
This test verifies:
- Proper app model config to workflow conversion
- Correct workflow graph structure
- Proper node creation and configuration
- Database state management
"""
# Arrange: Create test data
fake = Faker()
account, tenant = self._create_test_account_and_tenant(
db_session_with_containers, mock_external_service_dependencies
)
app = self._create_test_app(db_session_with_containers, mock_external_service_dependencies, tenant, account)
# Act: Execute the conversion
workflow_converter = WorkflowConverter()
workflow = workflow_converter.convert_app_model_config_to_workflow(
app_model=app,
app_model_config=app.app_model_config,
account_id=account.id,
)
# Assert: Verify the expected outcomes
assert workflow is not None
assert workflow.tenant_id == app.tenant_id
assert workflow.app_id == app.id
assert workflow.type == "chat"
assert workflow.version == Workflow.VERSION_DRAFT
assert workflow.created_by == account.id
# Verify workflow graph structure
graph = json.loads(workflow.graph)
assert "nodes" in graph
assert "edges" in graph
assert len(graph["nodes"]) > 0
assert len(graph["edges"]) > 0
# Verify start node exists
start_node = next((node for node in graph["nodes"] if node["data"]["type"] == "start"), None)
assert start_node is not None
assert start_node["id"] == "start"
# Verify LLM node exists
llm_node = next((node for node in graph["nodes"] if node["data"]["type"] == "llm"), None)
assert llm_node is not None
assert llm_node["id"] == "llm"
# Verify answer node exists for chat mode
answer_node = next((node for node in graph["nodes"] if node["data"]["type"] == "answer"), None)
assert answer_node is not None
assert answer_node["id"] == "answer"
# Verify database state
from extensions.ext_database import db
db.session.refresh(workflow)
assert workflow.id is not None
# Verify features were set
features = json.loads(workflow._features) if workflow._features else {}
assert isinstance(features, dict)
def test_convert_to_start_node_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful conversion to start node.
This test verifies:
- Proper start node creation with variables
- Correct node structure and data
- Variable encoding and formatting
"""
# Arrange: Create test variables
variables = [
VariableEntity(
variable="text_input",
label="Text Input",
type=VariableEntityType.TEXT_INPUT,
),
VariableEntity(
variable="number_input",
label="Number Input",
type=VariableEntityType.NUMBER,
),
]
# Act: Execute the conversion
workflow_converter = WorkflowConverter()
start_node = workflow_converter._convert_to_start_node(variables=variables)
# Assert: Verify the expected outcomes
assert start_node is not None
assert start_node["id"] == "start"
assert start_node["data"]["title"] == "START"
assert start_node["data"]["type"] == "start"
assert len(start_node["data"]["variables"]) == 2
# Verify variable encoding
first_variable = start_node["data"]["variables"][0]
assert first_variable["variable"] == "text_input"
assert first_variable["label"] == "Text Input"
assert first_variable["type"] == "text-input"
second_variable = start_node["data"]["variables"][1]
assert second_variable["variable"] == "number_input"
assert second_variable["label"] == "Number Input"
assert second_variable["type"] == "number"
def test_convert_to_http_request_node_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful conversion to HTTP request node.
This test verifies:
- Proper HTTP request node creation
- Correct API configuration and authorization
- Code node creation for response parsing
- External data variable mapping
"""
# Arrange: Create test data
fake = Faker()
account, tenant = self._create_test_account_and_tenant(
db_session_with_containers, mock_external_service_dependencies
)
app = self._create_test_app(db_session_with_containers, mock_external_service_dependencies, tenant, account)
# Create API based extension
api_based_extension = APIBasedExtension(
tenant_id=tenant.id,
name="Test API Extension",
api_key="encrypted_api_key",
api_endpoint="https://api.example.com/test",
)
from extensions.ext_database import db
db.session.add(api_based_extension)
db.session.commit()
# Mock encrypter
mock_external_service_dependencies["encrypter"].decrypt_token.return_value = "decrypted_api_key"
variables = [
VariableEntity(
variable="user_input",
label="User Input",
type=VariableEntityType.TEXT_INPUT,
)
]
external_data_variables = [
ExternalDataVariableEntity(
variable="external_data", type="api", config={"api_based_extension_id": api_based_extension.id}
)
]
# Act: Execute the conversion
workflow_converter = WorkflowConverter()
nodes, external_data_variable_node_mapping = workflow_converter._convert_to_http_request_node(
app_model=app,
variables=variables,
external_data_variables=external_data_variables,
)
# Assert: Verify the expected outcomes
assert len(nodes) == 2 # HTTP request node + code node
assert len(external_data_variable_node_mapping) == 1
# Verify HTTP request node
http_request_node = nodes[0]
assert http_request_node["data"]["type"] == "http-request"
assert http_request_node["data"]["method"] == "post"
assert http_request_node["data"]["url"] == api_based_extension.api_endpoint
assert http_request_node["data"]["authorization"]["type"] == "api-key"
assert http_request_node["data"]["authorization"]["config"]["type"] == "bearer"
assert http_request_node["data"]["authorization"]["config"]["api_key"] == "decrypted_api_key"
# Verify code node
code_node = nodes[1]
assert code_node["data"]["type"] == "code"
assert code_node["data"]["code_language"] == "python3"
assert "response_json" in code_node["data"]["variables"][0]["variable"]
# Verify mapping
assert external_data_variable_node_mapping["external_data"] == code_node["id"]
def test_convert_to_knowledge_retrieval_node_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful conversion to knowledge retrieval node.
This test verifies:
- Proper knowledge retrieval node creation
- Correct dataset configuration
- Model configuration integration
- Query variable selector setup
"""
# Arrange: Create test data
fake = Faker()
account, tenant = self._create_test_account_and_tenant(
db_session_with_containers, mock_external_service_dependencies
)
# Create dataset config
dataset_config = DatasetEntity(
dataset_ids=["dataset_1", "dataset_2"],
retrieve_config=DatasetRetrieveConfigEntity(
retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE,
top_k=10,
score_threshold=0.8,
reranking_model={"provider": "cohere", "model": "rerank-v2"},
reranking_enabled=True,
),
)
model_config = ModelConfigEntity(
provider="openai",
model="gpt-4",
mode=LLMMode.CHAT.value,
parameters={"temperature": 0.7},
stop=[],
)
# Act: Execute the conversion for advanced chat mode
workflow_converter = WorkflowConverter()
node = workflow_converter._convert_to_knowledge_retrieval_node(
new_app_mode=AppMode.ADVANCED_CHAT,
dataset_config=dataset_config,
model_config=model_config,
)
# Assert: Verify the expected outcomes
assert node is not None
assert node["data"]["type"] == "knowledge-retrieval"
assert node["data"]["title"] == "KNOWLEDGE RETRIEVAL"
assert node["data"]["dataset_ids"] == ["dataset_1", "dataset_2"]
assert node["data"]["retrieval_mode"] == "multiple"
assert node["data"]["query_variable_selector"] == ["sys", "query"]
# Verify multiple retrieval config
multiple_config = node["data"]["multiple_retrieval_config"]
assert multiple_config["top_k"] == 10
assert multiple_config["score_threshold"] == 0.8
assert multiple_config["reranking_model"]["provider"] == "cohere"
assert multiple_config["reranking_model"]["model"] == "rerank-v2"
# Verify single retrieval config is None for multiple strategy
assert node["data"]["single_retrieval_config"] is None

0
api/tests/test_containers_integration_tests/tasks/__init__.py Normal file
View File

786
api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py Normal file
View File

@ -0,0 +1,786 @@
from unittest.mock import MagicMock, patch
import pytest
from faker import Faker
from core.rag.index_processor.constant.index_type import IndexType
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, DatasetAutoDisableLog, Document, DocumentSegment
from tasks.add_document_to_index_task import add_document_to_index_task
class TestAddDocumentToIndexTask:
"""Integration tests for add_document_to_index_task using testcontainers."""
@pytest.fixture
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,
):
# Setup mock index processor
mock_processor = MagicMock()
mock_index_processor_factory.return_value.init_index_processor.return_value = mock_processor
yield {
"index_processor_factory": mock_index_processor_factory,
"index_processor": mock_processor,
}
def _create_test_dataset_and_document(self, db_session_with_containers, mock_external_service_dependencies):
"""
Helper method to create a test dataset and document for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
mock_external_service_dependencies: Mock dependencies
Returns:
tuple: (dataset, document) - Created dataset and document instances
"""
fake = Faker()
# Create account and tenant
account = Account(
email=fake.email(),
name=fake.name(),
interface_language="en-US",
status="active",
)
db.session.add(account)
db.session.commit()
tenant = Tenant(
name=fake.company(),
status="normal",
)
db.session.add(tenant)
db.session.commit()
# Create tenant-account join
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER.value,
current=True,
)
db.session.add(join)
db.session.commit()
# Create dataset
dataset = Dataset(
id=fake.uuid4(),
tenant_id=tenant.id,
name=fake.company(),
description=fake.text(max_nb_chars=100),
data_source_type="upload_file",
indexing_technique="high_quality",
created_by=account.id,
)
db.session.add(dataset)
db.session.commit()
# Create document
document = Document(
id=fake.uuid4(),
tenant_id=tenant.id,
dataset_id=dataset.id,
position=1,
data_source_type="upload_file",
batch="test_batch",
name=fake.file_name(),
created_from="upload_file",
created_by=account.id,
indexing_status="completed",
enabled=True,
doc_form=IndexType.PARAGRAPH_INDEX,
)
db.session.add(document)
db.session.commit()
# Refresh dataset to ensure doc_form property works correctly
db.session.refresh(dataset)
return dataset, document
def _create_test_segments(self, db_session_with_containers, document, dataset):
"""
Helper method to create test document segments.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
document: Document instance
dataset: Dataset instance
Returns:
list: List of created DocumentSegment instances
"""
fake = Faker()
segments = []
for i in range(3):
segment = DocumentSegment(
id=fake.uuid4(),
tenant_id=document.tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=i,
content=fake.text(max_nb_chars=200),
word_count=len(fake.text(max_nb_chars=200).split()),
tokens=len(fake.text(max_nb_chars=200).split()) * 2,
index_node_id=f"node_{i}",
index_node_hash=f"hash_{i}",
enabled=False,
status="completed",
created_by=document.created_by,
)
db.session.add(segment)
segments.append(segment)
db.session.commit()
return segments
def test_add_document_to_index_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful document indexing with paragraph index type.
This test verifies:
- Proper document retrieval from database
- Correct segment processing and document creation
- Index processor integration
- Database state updates
- Segment status changes
- Redis cache key deletion
"""
# Arrange: Create test data
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
segments = self._create_test_segments(db_session_with_containers, document, dataset)
# Set up Redis cache key to simulate indexing in progress
indexing_cache_key = f"document_{document.id}_indexing"
redis_client.set(indexing_cache_key, "processing", ex=300) # 5 minutes expiry
# Verify cache key exists
assert redis_client.exists(indexing_cache_key) == 1
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify the expected outcomes
# Verify index processor was called correctly
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(IndexType.PARAGRAPH_INDEX)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
# Verify database state changes
db.session.refresh(document)
for segment in segments:
db.session.refresh(segment)
assert segment.enabled is True
assert segment.disabled_at is None
assert segment.disabled_by is None
# Verify Redis cache key was deleted
assert redis_client.exists(indexing_cache_key) == 0
def test_add_document_to_index_with_different_index_type(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test document indexing with different index types.
This test verifies:
- Proper handling of different index types
- Index processor factory integration
- Document processing with various configurations
- Redis cache key deletion
"""
# Arrange: Create test data with different index type
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
# Update document to use different index type
document.doc_form = IndexType.QA_INDEX
db.session.commit()
# Refresh dataset to ensure doc_form property reflects the updated document
db.session.refresh(dataset)
# Create segments
segments = self._create_test_segments(db_session_with_containers, document, dataset)
# Set up Redis cache key
indexing_cache_key = f"document_{document.id}_indexing"
redis_client.set(indexing_cache_key, "processing", ex=300)
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify different index type handling
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(IndexType.QA_INDEX)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
# Verify the load method was called with correct parameters
call_args = mock_external_service_dependencies["index_processor"].load.call_args
assert call_args is not None
documents = call_args[0][1] # Second argument should be documents list
assert len(documents) == 3
# Verify database state changes
db.session.refresh(document)
for segment in segments:
db.session.refresh(segment)
assert segment.enabled is True
assert segment.disabled_at is None
assert segment.disabled_by is None
# Verify Redis cache key was deleted
assert redis_client.exists(indexing_cache_key) == 0
def test_add_document_to_index_document_not_found(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test handling of non-existent document.
This test verifies:
- Proper error handling for missing documents
- Early return without processing
- Database session cleanup
- No unnecessary index processor calls
- Redis cache key not affected (since it was never created)
"""
# Arrange: Use non-existent document ID
fake = Faker()
non_existent_id = fake.uuid4()
# Act: Execute the task with non-existent document
add_document_to_index_task(non_existent_id)
# Assert: Verify no processing occurred
mock_external_service_dependencies["index_processor_factory"].assert_not_called()
mock_external_service_dependencies["index_processor"].load.assert_not_called()
# Note: redis_client.delete is not called when document is not found
# because indexing_cache_key is not defined in that case
def test_add_document_to_index_invalid_indexing_status(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test handling of document with invalid indexing status.
This test verifies:
- Early return when indexing_status is not "completed"
- No index processing for documents not ready for indexing
- Proper database session cleanup
- No unnecessary external service calls
- Redis cache key not affected
"""
# Arrange: Create test data with invalid indexing status
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
# Set invalid indexing status
document.indexing_status = "processing"
db.session.commit()
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify no processing occurred
mock_external_service_dependencies["index_processor_factory"].assert_not_called()
mock_external_service_dependencies["index_processor"].load.assert_not_called()
def test_add_document_to_index_dataset_not_found(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test handling when document's dataset doesn't exist.
This test verifies:
- Proper error handling when dataset is missing
- Document status is set to error
- Document is disabled
- Error information is recorded
- Redis cache is cleared despite error
"""
# Arrange: Create test data
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
# Set up Redis cache key
indexing_cache_key = f"document_{document.id}_indexing"
redis_client.set(indexing_cache_key, "processing", ex=300)
# Delete the dataset to simulate dataset not found scenario
db.session.delete(dataset)
db.session.commit()
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify error handling
db.session.refresh(document)
assert document.enabled is False
assert document.indexing_status == "error"
assert document.error is not None
assert "doesn't exist" in document.error
assert document.disabled_at is not None
# Verify no index processing occurred
mock_external_service_dependencies["index_processor_factory"].assert_not_called()
mock_external_service_dependencies["index_processor"].load.assert_not_called()
# Verify redis cache was cleared despite error
assert redis_client.exists(indexing_cache_key) == 0
def test_add_document_to_index_with_parent_child_structure(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test document indexing with parent-child structure.
This test verifies:
- Proper handling of PARENT_CHILD_INDEX type
- Child document creation from segments
- Correct document structure for parent-child indexing
- Index processor receives properly structured documents
- Redis cache key deletion
"""
# Arrange: Create test data with parent-child index type
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
# Update document to use parent-child index type
document.doc_form = IndexType.PARENT_CHILD_INDEX
db.session.commit()
# Refresh dataset to ensure doc_form property reflects the updated document
db.session.refresh(dataset)
# Create segments with mock child chunks
segments = self._create_test_segments(db_session_with_containers, document, dataset)
# Set up Redis cache key
indexing_cache_key = f"document_{document.id}_indexing"
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:
# Setup mock to return child chunks for each segment
mock_child_chunks = []
for i in range(2): # Each segment has 2 child chunks
mock_child = MagicMock()
mock_child.content = f"child_content_{i}"
mock_child.index_node_id = f"child_node_{i}"
mock_child.index_node_hash = f"child_hash_{i}"
mock_child_chunks.append(mock_child)
mock_get_child_chunks.return_value = mock_child_chunks
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify parent-child index processing
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(
IndexType.PARENT_CHILD_INDEX
)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
# Verify the load method was called with correct parameters
call_args = mock_external_service_dependencies["index_processor"].load.call_args
assert call_args is not None
documents = call_args[0][1] # Second argument should be documents list
assert len(documents) == 3 # 3 segments
# Verify each document has children
for doc in documents:
assert hasattr(doc, "children")
assert len(doc.children) == 2 # Each document has 2 children
# Verify database state changes
db.session.refresh(document)
for segment in segments:
db.session.refresh(segment)
assert segment.enabled is True
assert segment.disabled_at is None
assert segment.disabled_by is None
# Verify redis cache was cleared
assert redis_client.exists(indexing_cache_key) == 0
def test_add_document_to_index_with_no_segments_to_process(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test document indexing when no segments need processing.
This test verifies:
- Proper handling when all segments are already enabled
- Index processing still occurs but with empty documents list
- Auto disable log deletion still occurs
- Redis cache is cleared
"""
# Arrange: Create test data
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
# Create segments that are already enabled
fake = Faker()
segments = []
for i in range(3):
segment = DocumentSegment(
id=fake.uuid4(),
tenant_id=document.tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=i,
content=fake.text(max_nb_chars=200),
word_count=len(fake.text(max_nb_chars=200).split()),
tokens=len(fake.text(max_nb_chars=200).split()) * 2,
index_node_id=f"node_{i}",
index_node_hash=f"hash_{i}",
enabled=True, # Already enabled
status="completed",
created_by=document.created_by,
)
db.session.add(segment)
segments.append(segment)
db.session.commit()
# Set up Redis cache key
indexing_cache_key = f"document_{document.id}_indexing"
redis_client.set(indexing_cache_key, "processing", ex=300)
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify index processing occurred but with empty documents list
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(IndexType.PARAGRAPH_INDEX)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
# Verify the load method was called with empty documents list
call_args = mock_external_service_dependencies["index_processor"].load.call_args
assert call_args is not None
documents = call_args[0][1] # Second argument should be documents list
assert len(documents) == 0 # No segments to process
# Verify redis cache was cleared
assert redis_client.exists(indexing_cache_key) == 0
def test_add_document_to_index_auto_disable_log_deletion(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test that auto disable logs are properly deleted during indexing.
This test verifies:
- Auto disable log entries are deleted for the document
- Database state is properly managed
- Index processing continues normally
- Redis cache key deletion
"""
# Arrange: Create test data
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
segments = self._create_test_segments(db_session_with_containers, document, dataset)
# Create some auto disable log entries
fake = Faker()
auto_disable_logs = []
for i in range(2):
log_entry = DatasetAutoDisableLog(
id=fake.uuid4(),
tenant_id=document.tenant_id,
dataset_id=dataset.id,
document_id=document.id,
)
db.session.add(log_entry)
auto_disable_logs.append(log_entry)
db.session.commit()
# Set up Redis cache key
indexing_cache_key = f"document_{document.id}_indexing"
redis_client.set(indexing_cache_key, "processing", ex=300)
# Verify logs exist before processing
existing_logs = (
db.session.query(DatasetAutoDisableLog).where(DatasetAutoDisableLog.document_id == document.id).all()
)
assert len(existing_logs) == 2
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify auto disable logs were deleted
remaining_logs = (
db.session.query(DatasetAutoDisableLog).where(DatasetAutoDisableLog.document_id == document.id).all()
)
assert len(remaining_logs) == 0
# Verify index processing occurred normally
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(IndexType.PARAGRAPH_INDEX)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
# Verify segments were enabled
for segment in segments:
db.session.refresh(segment)
assert segment.enabled is True
# Verify redis cache was cleared
assert redis_client.exists(indexing_cache_key) == 0
def test_add_document_to_index_general_exception_handling(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test general exception handling during indexing process.
This test verifies:
- Exceptions are properly caught and handled
- Document status is set to error
- Document is disabled
- Error information is recorded
- Redis cache is still cleared
- Database session is properly closed
"""
# Arrange: Create test data
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
segments = self._create_test_segments(db_session_with_containers, document, dataset)
# Set up Redis cache key
indexing_cache_key = f"document_{document.id}_indexing"
redis_client.set(indexing_cache_key, "processing", ex=300)
# Mock the index processor to raise an exception
mock_external_service_dependencies["index_processor"].load.side_effect = Exception("Index processing failed")
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify error handling
db.session.refresh(document)
assert document.enabled is False
assert document.indexing_status == "error"
assert document.error is not None
assert "Index processing failed" in document.error
assert document.disabled_at is not None
# Verify segments were not enabled due to error
for segment in segments:
db.session.refresh(segment)
assert segment.enabled is False # Should remain disabled due to error
# Verify redis cache was still cleared despite error
assert redis_client.exists(indexing_cache_key) == 0
def test_add_document_to_index_segment_filtering_edge_cases(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test segment filtering with various edge cases.
This test verifies:
- Only segments with enabled=False and status="completed" are processed
- Segments are ordered by position correctly
- Mixed segment states are handled properly
- Redis cache key deletion
"""
# Arrange: Create test data
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
# Create segments with mixed states
fake = Faker()
segments = []
# Segment 1: Should be processed (enabled=False, status="completed")
segment1 = DocumentSegment(
id=fake.uuid4(),
tenant_id=document.tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=0,
content=fake.text(max_nb_chars=200),
word_count=len(fake.text(max_nb_chars=200).split()),
tokens=len(fake.text(max_nb_chars=200).split()) * 2,
index_node_id="node_0",
index_node_hash="hash_0",
enabled=False,
status="completed",
created_by=document.created_by,
)
db.session.add(segment1)
segments.append(segment1)
# Segment 2: Should NOT be processed (enabled=True, status="completed")
segment2 = DocumentSegment(
id=fake.uuid4(),
tenant_id=document.tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content=fake.text(max_nb_chars=200),
word_count=len(fake.text(max_nb_chars=200).split()),
tokens=len(fake.text(max_nb_chars=200).split()) * 2,
index_node_id="node_1",
index_node_hash="hash_1",
enabled=True, # Already enabled
status="completed",
created_by=document.created_by,
)
db.session.add(segment2)
segments.append(segment2)
# Segment 3: Should NOT be processed (enabled=False, status="processing")
segment3 = DocumentSegment(
id=fake.uuid4(),
tenant_id=document.tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=2,
content=fake.text(max_nb_chars=200),
word_count=len(fake.text(max_nb_chars=200).split()),
tokens=len(fake.text(max_nb_chars=200).split()) * 2,
index_node_id="node_2",
index_node_hash="hash_2",
enabled=False,
status="processing", # Not completed
created_by=document.created_by,
)
db.session.add(segment3)
segments.append(segment3)
# Segment 4: Should be processed (enabled=False, status="completed")
segment4 = DocumentSegment(
id=fake.uuid4(),
tenant_id=document.tenant_id,
dataset_id=dataset.id,
document_id=document.id,
position=3,
content=fake.text(max_nb_chars=200),
word_count=len(fake.text(max_nb_chars=200).split()),
tokens=len(fake.text(max_nb_chars=200).split()) * 2,
index_node_id="node_3",
index_node_hash="hash_3",
enabled=False,
status="completed",
created_by=document.created_by,
)
db.session.add(segment4)
segments.append(segment4)
db.session.commit()
# Set up Redis cache key
indexing_cache_key = f"document_{document.id}_indexing"
redis_client.set(indexing_cache_key, "processing", ex=300)
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify only eligible segments were processed
mock_external_service_dependencies["index_processor_factory"].assert_called_once_with(IndexType.PARAGRAPH_INDEX)
mock_external_service_dependencies["index_processor"].load.assert_called_once()
# Verify the load method was called with correct parameters
call_args = mock_external_service_dependencies["index_processor"].load.call_args
assert call_args is not None
documents = call_args[0][1] # Second argument should be documents list
assert len(documents) == 2 # Only 2 segments should be processed
# Verify correct segments were processed (by position order)
assert documents[0].metadata["doc_id"] == "node_0" # position 0
assert documents[1].metadata["doc_id"] == "node_3" # position 3
# Verify database state changes
db.session.refresh(document)
db.session.refresh(segment1)
db.session.refresh(segment2)
db.session.refresh(segment3)
db.session.refresh(segment4)
# All segments should be enabled because the task updates ALL segments for the document
assert segment1.enabled is True
assert segment2.enabled is True # Was already enabled, now updated to True
assert segment3.enabled is True # Was not processed but still updated to True
assert segment4.enabled is True
# Verify redis cache was cleared
assert redis_client.exists(indexing_cache_key) == 0
def test_add_document_to_index_comprehensive_error_scenarios(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test comprehensive error scenarios and recovery.
This test verifies:
- Multiple types of exceptions are handled properly
- Error state is consistently managed
- Resource cleanup occurs in all error cases
- Database session management is robust
- Redis cache key deletion in all scenarios
"""
# Arrange: Create test data
dataset, document = self._create_test_dataset_and_document(
db_session_with_containers, mock_external_service_dependencies
)
segments = self._create_test_segments(db_session_with_containers, document, dataset)
# Test different exception types
test_exceptions = [
("Database connection error", Exception("Database connection failed")),
("Index processor error", RuntimeError("Index processor initialization failed")),
("Memory error", MemoryError("Out of memory")),
("Value error", ValueError("Invalid index type")),
]
for error_name, exception in test_exceptions:
# Reset mocks for each test
mock_external_service_dependencies["index_processor"].load.side_effect = exception
# Reset document state
document.enabled = True
document.indexing_status = "completed"
document.error = None
document.disabled_at = None
db.session.commit()
# Set up Redis cache key
indexing_cache_key = f"document_{document.id}_indexing"
redis_client.set(indexing_cache_key, "processing", ex=300)
# Act: Execute the task
add_document_to_index_task(document.id)
# Assert: Verify consistent error handling
db.session.refresh(document)
assert document.enabled is False, f"Document should be disabled for {error_name}"
assert document.indexing_status == "error", f"Document status should be error for {error_name}"
assert document.error is not None, f"Error should be recorded for {error_name}"
assert str(exception) in document.error, f"Error message should contain exception for {error_name}"
assert document.disabled_at is not None, f"Disabled timestamp should be set for {error_name}"
# Verify segments remain disabled due to error
for segment in segments:
db.session.refresh(segment)
assert segment.enabled is False, f"Segments should remain disabled for {error_name}"
# Verify redis cache was still cleared despite error
assert redis_client.exists(indexing_cache_key) == 0, f"Redis cache should be cleared for {error_name}"

720
api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py Normal file
View File

@ -0,0 +1,720 @@
"""
Integration tests for batch_clean_document_task using testcontainers.
This module tests the batch document cleaning functionality with real database
and storage containers to ensure proper cleanup of documents, segments, and files.
"""
import json
import uuid
from unittest.mock import Mock, patch
import pytest
from faker import Faker
from extensions.ext_database import db
from libs.datetime_utils import naive_utc_now
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, Document, DocumentSegment
from models.model import UploadFile
from tasks.batch_clean_document_task import batch_clean_document_task
class TestBatchCleanDocumentTask:
"""Integration tests for batch_clean_document_task using testcontainers."""
@pytest.fixture
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("extensions.ext_storage.storage") as mock_storage,
patch("core.rag.index_processor.index_processor_factory.IndexProcessorFactory") as mock_index_factory,
patch("core.tools.utils.web_reader_tool.get_image_upload_file_ids") as mock_get_image_ids,
):
# Setup default mock returns
mock_storage.delete.return_value = None
# Mock index processor
mock_index_processor = Mock()
mock_index_processor.clean.return_value = None
mock_index_factory.return_value.init_index_processor.return_value = mock_index_processor
# Mock image file ID extraction
mock_get_image_ids.return_value = []
yield {
"storage": mock_storage,
"index_factory": mock_index_factory,
"index_processor": mock_index_processor,
"get_image_ids": mock_get_image_ids,
}
def _create_test_account(self, db_session_with_containers):
"""
Helper method to create a test account for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
Returns:
Account: Created account instance
"""
fake = Faker()
# Create account
account = Account(
email=fake.email(),
name=fake.name(),
interface_language="en-US",
status="active",
)
db.session.add(account)
db.session.commit()
# Create tenant for the account
tenant = Tenant(
name=fake.company(),
status="normal",
)
db.session.add(tenant)
db.session.commit()
# Create tenant-account join
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER.value,
current=True,
)
db.session.add(join)
db.session.commit()
# Set current tenant for account
account.current_tenant = tenant
return account
def _create_test_dataset(self, db_session_with_containers, account):
"""
Helper method to create a test dataset for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
Returns:
Dataset: Created dataset instance
"""
fake = Faker()
dataset = Dataset(
id=str(uuid.uuid4()),
tenant_id=account.current_tenant.id,
name=fake.word(),
description=fake.sentence(),
data_source_type="upload_file",
created_by=account.id,
embedding_model="text-embedding-ada-002",
embedding_model_provider="openai",
)
db.session.add(dataset)
db.session.commit()
return dataset
def _create_test_document(self, db_session_with_containers, dataset, account):
"""
Helper method to create a test document for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
dataset: Dataset instance
account: Account instance
Returns:
Document: Created document instance
"""
fake = Faker()
document = Document(
id=str(uuid.uuid4()),
tenant_id=account.current_tenant.id,
dataset_id=dataset.id,
position=0,
name=fake.word(),
data_source_type="upload_file",
data_source_info=json.dumps({"upload_file_id": str(uuid.uuid4())}),
batch="test_batch",
created_from="test",
created_by=account.id,
indexing_status="completed",
doc_form="text_model",
)
db.session.add(document)
db.session.commit()
return document
def _create_test_document_segment(self, db_session_with_containers, document, account):
"""
Helper method to create a test document segment for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
document: Document instance
account: Account instance
Returns:
DocumentSegment: Created document segment instance
"""
fake = Faker()
segment = DocumentSegment(
id=str(uuid.uuid4()),
tenant_id=account.current_tenant.id,
dataset_id=document.dataset_id,
document_id=document.id,
position=0,
content=fake.text(),
word_count=100,
tokens=50,
index_node_id=str(uuid.uuid4()),
created_by=account.id,
status="completed",
)
db.session.add(segment)
db.session.commit()
return segment
def _create_test_upload_file(self, db_session_with_containers, account):
"""
Helper method to create a test upload file for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
Returns:
UploadFile: Created upload file instance
"""
fake = Faker()
from models.enums import CreatorUserRole
upload_file = UploadFile(
tenant_id=account.current_tenant.id,
storage_type="local",
key=f"test_files/{fake.file_name()}",
name=fake.file_name(),
size=1024,
extension="txt",
mime_type="text/plain",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=naive_utc_now(),
used=False,
)
db.session.add(upload_file)
db.session.commit()
return upload_file
def test_batch_clean_document_task_successful_cleanup(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful cleanup of documents with segments and files.
This test verifies that the task properly cleans up:
- Document segments from the index
- Associated image files from storage
- Upload files from storage and database
"""
# Create test data
account = self._create_test_account(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account)
document = self._create_test_document(db_session_with_containers, dataset, account)
segment = self._create_test_document_segment(db_session_with_containers, document, account)
upload_file = self._create_test_upload_file(db_session_with_containers, account)
# Update document to reference the upload file
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
db.session.commit()
# Store original IDs for verification
document_id = document.id
segment_id = segment.id
file_id = upload_file.id
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[file_id]
)
# Verify that the task completed successfully
# The task should have processed the segment and cleaned up the database
# Verify database cleanup
db.session.commit() # Ensure all changes are committed
# Check that segment is deleted
deleted_segment = db.session.query(DocumentSegment).filter_by(id=segment_id).first()
assert deleted_segment is None
# Check that upload file is deleted
deleted_file = db.session.query(UploadFile).filter_by(id=file_id).first()
assert deleted_file is None
def test_batch_clean_document_task_with_image_files(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test cleanup of documents containing image references.
This test verifies that the task properly handles documents with
image content and cleans up associated segments.
"""
# Create test data
account = self._create_test_account(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account)
document = self._create_test_document(db_session_with_containers, dataset, account)
# Create segment with simple content (no image references)
segment = DocumentSegment(
id=str(uuid.uuid4()),
tenant_id=account.current_tenant.id,
dataset_id=document.dataset_id,
document_id=document.id,
position=0,
content="Simple text content without images",
word_count=100,
tokens=50,
index_node_id=str(uuid.uuid4()),
created_by=account.id,
status="completed",
)
db.session.add(segment)
db.session.commit()
# Store original IDs for verification
segment_id = segment.id
document_id = document.id
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[]
)
# Verify database cleanup
db.session.commit()
# Check that segment is deleted
deleted_segment = db.session.query(DocumentSegment).filter_by(id=segment_id).first()
assert deleted_segment is None
# Verify that the task completed successfully by checking the log output
# The task should have processed the segment and cleaned up the database
def test_batch_clean_document_task_no_segments(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test cleanup when document has no segments.
This test verifies that the task handles documents without segments
gracefully and still cleans up associated files.
"""
# Create test data without segments
account = self._create_test_account(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account)
document = self._create_test_document(db_session_with_containers, dataset, account)
upload_file = self._create_test_upload_file(db_session_with_containers, account)
# Update document to reference the upload file
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
db.session.commit()
# Store original IDs for verification
document_id = document.id
file_id = upload_file.id
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[file_id]
)
# Verify that the task completed successfully
# Since there are no segments, the task should handle this gracefully
# Verify database cleanup
db.session.commit()
# Check that upload file is deleted
deleted_file = db.session.query(UploadFile).filter_by(id=file_id).first()
assert deleted_file is None
# Verify database cleanup
db.session.commit()
# Check that upload file is deleted
deleted_file = db.session.query(UploadFile).filter_by(id=file_id).first()
assert deleted_file is None
def test_batch_clean_document_task_dataset_not_found(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test cleanup when dataset is not found.
This test verifies that the task properly handles the case where
the specified dataset does not exist in the database.
"""
# Create test data
account = self._create_test_account(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account)
document = self._create_test_document(db_session_with_containers, dataset, account)
# Store original IDs for verification
document_id = document.id
dataset_id = dataset.id
# Delete the dataset to simulate not found scenario
db.session.delete(dataset)
db.session.commit()
# Execute the task with non-existent dataset
batch_clean_document_task(document_ids=[document_id], dataset_id=dataset_id, doc_form="text_model", file_ids=[])
# Verify that no index processing occurred
mock_external_service_dependencies["index_processor"].clean.assert_not_called()
# Verify that no storage operations occurred
mock_external_service_dependencies["storage"].delete.assert_not_called()
# Verify that no database cleanup occurred
db.session.commit()
# Document should still exist since cleanup failed
existing_document = db.session.query(Document).filter_by(id=document_id).first()
assert existing_document is not None
def test_batch_clean_document_task_storage_cleanup_failure(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test cleanup when storage operations fail.
This test verifies that the task continues processing even when
storage cleanup operations fail, ensuring database cleanup still occurs.
"""
# Create test data
account = self._create_test_account(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account)
document = self._create_test_document(db_session_with_containers, dataset, account)
segment = self._create_test_document_segment(db_session_with_containers, document, account)
upload_file = self._create_test_upload_file(db_session_with_containers, account)
# Update document to reference the upload file
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
db.session.commit()
# Store original IDs for verification
document_id = document.id
segment_id = segment.id
file_id = upload_file.id
# Mock storage.delete to raise an exception
mock_external_service_dependencies["storage"].delete.side_effect = Exception("Storage error")
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[file_id]
)
# Verify that the task completed successfully despite storage failure
# The task should continue processing even when storage operations fail
# Verify database cleanup still occurred despite storage failure
db.session.commit()
# Check that segment is deleted from database
deleted_segment = db.session.query(DocumentSegment).filter_by(id=segment_id).first()
assert deleted_segment is None
# Check that upload file is deleted from database
deleted_file = db.session.query(UploadFile).filter_by(id=file_id).first()
assert deleted_file is None
def test_batch_clean_document_task_multiple_documents(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test cleanup of multiple documents in a single batch operation.
This test verifies that the task can handle multiple documents
efficiently and cleans up all associated resources.
"""
# Create test data for multiple documents
account = self._create_test_account(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account)
documents = []
segments = []
upload_files = []
# Create 3 documents with segments and files
for i in range(3):
document = self._create_test_document(db_session_with_containers, dataset, account)
segment = self._create_test_document_segment(db_session_with_containers, document, account)
upload_file = self._create_test_upload_file(db_session_with_containers, account)
# Update document to reference the upload file
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
documents.append(document)
segments.append(segment)
upload_files.append(upload_file)
db.session.commit()
# Store original IDs for verification
document_ids = [doc.id for doc in documents]
segment_ids = [seg.id for seg in segments]
file_ids = [file.id for file in upload_files]
# Execute the task with multiple documents
batch_clean_document_task(
document_ids=document_ids, dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=file_ids
)
# Verify that the task completed successfully for all documents
# The task should process all documents and clean up all associated resources
# Verify database cleanup for all resources
db.session.commit()
# Check that all segments are deleted
for segment_id in segment_ids:
deleted_segment = db.session.query(DocumentSegment).filter_by(id=segment_id).first()
assert deleted_segment is None
# Check that all upload files are deleted
for file_id in file_ids:
deleted_file = db.session.query(UploadFile).filter_by(id=file_id).first()
assert deleted_file is None
def test_batch_clean_document_task_different_doc_forms(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test cleanup with different document form types.
This test verifies that the task properly handles different
document form types and creates the appropriate index processor.
"""
# Create test data
account = self._create_test_account(db_session_with_containers)
# Test different doc_form types
doc_forms = ["text_model", "qa_model", "hierarchical_model"]
for doc_form in doc_forms:
dataset = self._create_test_dataset(db_session_with_containers, account)
db.session.commit()
document = self._create_test_document(db_session_with_containers, dataset, account)
# Update document doc_form
document.doc_form = doc_form
db.session.commit()
segment = self._create_test_document_segment(db_session_with_containers, document, account)
# Store the ID before the object is deleted
segment_id = segment.id
try:
# Execute the task
batch_clean_document_task(
document_ids=[document.id], dataset_id=dataset.id, doc_form=doc_form, file_ids=[]
)
# Verify that the task completed successfully for this doc_form
# The task should handle different document forms correctly
# Verify database cleanup
db.session.commit()
# Check that segment is deleted
deleted_segment = db.session.query(DocumentSegment).filter_by(id=segment_id).first()
assert deleted_segment is None
except Exception as e:
# If the task fails due to external service issues (e.g., plugin daemon),
# we should still verify that the database state is consistent
# This is a common scenario in test environments where external services may not be available
db.session.commit()
# Check if the segment still exists (task may have failed before deletion)
existing_segment = db.session.query(DocumentSegment).filter_by(id=segment_id).first()
if existing_segment is not None:
# If segment still exists, the task failed before deletion
# This is acceptable in test environments with external service issues
pass
else:
# If segment was deleted, the task succeeded
pass
def test_batch_clean_document_task_large_batch_performance(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test cleanup performance with a large batch of documents.
This test verifies that the task can handle large batches efficiently
and maintains performance characteristics.
"""
import time
# Create test data for large batch
account = self._create_test_account(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account)
documents = []
segments = []
upload_files = []
# Create 10 documents with segments and files (larger batch)
batch_size = 10
for i in range(batch_size):
document = self._create_test_document(db_session_with_containers, dataset, account)
segment = self._create_test_document_segment(db_session_with_containers, document, account)
upload_file = self._create_test_upload_file(db_session_with_containers, account)
# Update document to reference the upload file
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
documents.append(document)
segments.append(segment)
upload_files.append(upload_file)
db.session.commit()
# Store original IDs for verification
document_ids = [doc.id for doc in documents]
segment_ids = [seg.id for seg in segments]
file_ids = [file.id for file in upload_files]
# Measure execution time
start_time = time.perf_counter()
# Execute the task with large batch
batch_clean_document_task(
document_ids=document_ids, dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=file_ids
)
end_time = time.perf_counter()
execution_time = end_time - start_time
# Verify performance characteristics (should complete within reasonable time)
assert execution_time < 5.0 # Should complete within 5 seconds
# Verify that the task completed successfully for the large batch
# The task should handle large batches efficiently
# Verify database cleanup for all resources
db.session.commit()
# Check that all segments are deleted
for segment_id in segment_ids:
deleted_segment = db.session.query(DocumentSegment).filter_by(id=segment_id).first()
assert deleted_segment is None
# Check that all upload files are deleted
for file_id in file_ids:
deleted_file = db.session.query(UploadFile).filter_by(id=file_id).first()
assert deleted_file is None
def test_batch_clean_document_task_integration_with_real_database(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test full integration with real database operations.
This test verifies that the task integrates properly with the
actual database and maintains data consistency throughout the process.
"""
# Create test data
account = self._create_test_account(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account)
# Create document with complex structure
document = self._create_test_document(db_session_with_containers, dataset, account)
# Create multiple segments for the document
segments = []
for i in range(3):
segment = DocumentSegment(
id=str(uuid.uuid4()),
tenant_id=account.current_tenant.id,
dataset_id=document.dataset_id,
document_id=document.id,
position=i,
content=f"Segment content {i} with some text",
word_count=50 + i * 10,
tokens=25 + i * 5,
index_node_id=str(uuid.uuid4()),
created_by=account.id,
status="completed",
)
segments.append(segment)
# Create upload file
upload_file = self._create_test_upload_file(db_session_with_containers, account)
# Update document to reference the upload file
document.data_source_info = json.dumps({"upload_file_id": upload_file.id})
# Add all to database
for segment in segments:
db.session.add(segment)
db.session.commit()
# Verify initial state
assert db.session.query(DocumentSegment).filter_by(document_id=document.id).count() == 3
assert db.session.query(UploadFile).filter_by(id=upload_file.id).first() is not None
# Store original IDs for verification
document_id = document.id
segment_ids = [seg.id for seg in segments]
file_id = upload_file.id
# Execute the task
batch_clean_document_task(
document_ids=[document_id], dataset_id=dataset.id, doc_form=dataset.doc_form, file_ids=[file_id]
)
# Verify that the task completed successfully
# The task should process all segments and clean up all associated resources
# Verify database cleanup
db.session.commit()
# Check that all segments are deleted
for segment_id in segment_ids:
deleted_segment = db.session.query(DocumentSegment).filter_by(id=segment_id).first()
assert deleted_segment is None
# Check that upload file is deleted
deleted_file = db.session.query(UploadFile).filter_by(id=file_id).first()
assert deleted_file is None
# Verify final database state
assert db.session.query(DocumentSegment).filter_by(document_id=document_id).count() == 0
assert db.session.query(UploadFile).filter_by(id=file_id).first() is None

737
api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py Normal file
View File

@ -0,0 +1,737 @@
"""
Integration tests for batch_create_segment_to_index_task using testcontainers.
This module provides comprehensive integration tests for the batch segment creation
and indexing task using TestContainers infrastructure. The tests ensure that the
task properly processes CSV files, creates document segments, and establishes
vector indexes in a real database environment.
All tests use the testcontainers infrastructure to ensure proper database isolation
and realistic testing scenarios with actual PostgreSQL and Redis instances.
"""
import uuid
from datetime import datetime
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from faker import Faker
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, Document, DocumentSegment
from models.enums import CreatorUserRole
from models.model import UploadFile
from tasks.batch_create_segment_to_index_task import batch_create_segment_to_index_task
class TestBatchCreateSegmentToIndexTask:
"""Integration tests for batch_create_segment_to_index_task using testcontainers."""
@pytest.fixture(autouse=True)
def cleanup_database(self, db_session_with_containers):
"""Clean up database before each test to ensure isolation."""
from extensions.ext_database import db
from extensions.ext_redis import redis_client
# Clear all test data
db.session.query(DocumentSegment).delete()
db.session.query(Document).delete()
db.session.query(Dataset).delete()
db.session.query(UploadFile).delete()
db.session.query(TenantAccountJoin).delete()
db.session.query(Tenant).delete()
db.session.query(Account).delete()
db.session.commit()
# Clear Redis cache
redis_client.flushdb()
@pytest.fixture
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,
):
# Setup default mock returns
mock_storage.download.return_value = None
# Mock embedding model for high quality indexing
mock_embedding_model = MagicMock()
mock_embedding_model.get_text_embedding_num_tokens.return_value = [10, 15, 20]
mock_model_manager_instance = MagicMock()
mock_model_manager_instance.get_model_instance.return_value = mock_embedding_model
mock_model_manager.return_value = mock_model_manager_instance
# Mock vector service
mock_vector_service.create_segments_vector.return_value = None
yield {
"storage": mock_storage,
"model_manager": mock_model_manager,
"vector_service": mock_vector_service,
"embedding_model": mock_embedding_model,
}
def _create_test_account_and_tenant(self, db_session_with_containers):
"""
Helper method to create a test account and tenant for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
Returns:
tuple: (Account, Tenant) created instances
"""
fake = Faker()
# Create account
account = Account(
email=fake.email(),
name=fake.name(),
interface_language="en-US",
status="active",
)
from extensions.ext_database import db
db.session.add(account)
db.session.commit()
# Create tenant for the account
tenant = Tenant(
name=fake.company(),
status="normal",
)
db.session.add(tenant)
db.session.commit()
# Create tenant-account join
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER.value,
current=True,
)
db.session.add(join)
db.session.commit()
# Set current tenant for account
account.current_tenant = tenant
return account, tenant
def _create_test_dataset(self, db_session_with_containers, account, tenant):
"""
Helper method to create a test dataset for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
tenant: Tenant instance
Returns:
Dataset: Created dataset instance
"""
fake = Faker()
dataset = Dataset(
tenant_id=tenant.id,
name=fake.company(),
description=fake.text(),
data_source_type="upload_file",
indexing_technique="high_quality",
embedding_model="text-embedding-ada-002",
embedding_model_provider="openai",
created_by=account.id,
)
from extensions.ext_database import db
db.session.add(dataset)
db.session.commit()
return dataset
def _create_test_document(self, db_session_with_containers, account, tenant, dataset):
"""
Helper method to create a test document for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
tenant: Tenant instance
dataset: Dataset instance
Returns:
Document: Created document instance
"""
fake = Faker()
document = Document(
tenant_id=tenant.id,
dataset_id=dataset.id,
position=1,
data_source_type="upload_file",
batch="test_batch",
name=fake.file_name(),
created_from="upload_file",
created_by=account.id,
indexing_status="completed",
enabled=True,
archived=False,
doc_form="text_model",
word_count=0,
)
from extensions.ext_database import db
db.session.add(document)
db.session.commit()
return document
def _create_test_upload_file(self, db_session_with_containers, account, tenant):
"""
Helper method to create a test upload file for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: Account instance
tenant: Tenant instance
Returns:
UploadFile: Created upload file instance
"""
fake = Faker()
upload_file = UploadFile(
tenant_id=tenant.id,
storage_type="local",
key=f"test_files/{fake.file_name()}",
name=fake.file_name(),
size=1024,
extension=".csv",
mime_type="text/csv",
created_by_role=CreatorUserRole.ACCOUNT,
created_by=account.id,
created_at=datetime.now(),
used=False,
)
from extensions.ext_database import db
db.session.add(upload_file)
db.session.commit()
return upload_file
def _create_test_csv_content(self, content_type="text_model"):
"""
Helper method to create test CSV content.
Args:
content_type: Type of content to create ("text_model" or "qa_model")
Returns:
str: CSV content as string
"""
if content_type == "qa_model":
csv_content = "content,answer\n"
csv_content += "This is the first segment content,This is the first answer\n"
csv_content += "This is the second segment content,This is the second answer\n"
csv_content += "This is the third segment content,This is the third answer\n"
else:
csv_content = "content\n"
csv_content += "This is the first segment content\n"
csv_content += "This is the second segment content\n"
csv_content += "This is the third segment content\n"
return csv_content
def test_batch_create_segment_to_index_task_success_text_model(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful batch creation of segments for text model documents.
This test verifies that the task can successfully:
1. Process a CSV file with text content
2. Create document segments with proper metadata
3. Update document word count
4. Create vector indexes
5. Set Redis cache status
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
# Create CSV content
csv_content = self._create_test_csv_content("text_model")
# Mock storage to return our CSV content
mock_storage = mock_external_service_dependencies["storage"]
def mock_download(key, file_path):
Path(file_path).write_text(csv_content, encoding="utf-8")
mock_storage.download.side_effect = mock_download
# Execute the task
job_id = str(uuid.uuid4())
batch_create_segment_to_index_task(
job_id=job_id,
upload_file_id=upload_file.id,
dataset_id=dataset.id,
document_id=document.id,
tenant_id=tenant.id,
user_id=account.id,
)
# Verify results
from extensions.ext_database import db
# Check that segments were created
segments = (
db.session.query(DocumentSegment)
.filter_by(document_id=document.id)
.order_by(DocumentSegment.position)
.all()
)
assert len(segments) == 3
# Verify segment content and metadata
for i, segment in enumerate(segments):
assert segment.tenant_id == tenant.id
assert segment.dataset_id == dataset.id
assert segment.document_id == document.id
assert segment.position == i + 1
assert segment.status == "completed"
assert segment.indexing_at is not None
assert segment.completed_at is not None
assert segment.answer is None # text_model doesn't have answers
# Check that document word count was updated
db.session.refresh(document)
assert document.word_count > 0
# Verify vector service was called
mock_vector_service = mock_external_service_dependencies["vector_service"]
mock_vector_service.create_segments_vector.assert_called_once()
# Check Redis cache was set
from extensions.ext_redis import redis_client
cache_key = f"segment_batch_import_{job_id}"
cache_value = redis_client.get(cache_key)
assert cache_value == b"completed"
def test_batch_create_segment_to_index_task_dataset_not_found(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test task failure when dataset does not exist.
This test verifies that the task properly handles error cases:
1. Fails gracefully when dataset is not found
2. Sets appropriate Redis cache status
3. Logs error information
4. Maintains database integrity
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
# Use non-existent IDs
non_existent_dataset_id = str(uuid.uuid4())
non_existent_document_id = str(uuid.uuid4())
# Execute the task with non-existent dataset
job_id = str(uuid.uuid4())
batch_create_segment_to_index_task(
job_id=job_id,
upload_file_id=upload_file.id,
dataset_id=non_existent_dataset_id,
document_id=non_existent_document_id,
tenant_id=tenant.id,
user_id=account.id,
)
# Verify error handling
# Check Redis cache was set to error status
from extensions.ext_redis import redis_client
cache_key = f"segment_batch_import_{job_id}"
cache_value = redis_client.get(cache_key)
assert cache_value == b"error"
# Verify no segments were created (since dataset doesn't exist)
from extensions.ext_database import db
segments = db.session.query(DocumentSegment).all()
assert len(segments) == 0
# Verify no documents were modified
documents = db.session.query(Document).all()
assert len(documents) == 0
def test_batch_create_segment_to_index_task_document_not_found(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test task failure when document does not exist.
This test verifies that the task properly handles error cases:
1. Fails gracefully when document is not found
2. Sets appropriate Redis cache status
3. Maintains database integrity
4. Logs appropriate error information
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
# Use non-existent document ID
non_existent_document_id = str(uuid.uuid4())
# Execute the task with non-existent document
job_id = str(uuid.uuid4())
batch_create_segment_to_index_task(
job_id=job_id,
upload_file_id=upload_file.id,
dataset_id=dataset.id,
document_id=non_existent_document_id,
tenant_id=tenant.id,
user_id=account.id,
)
# Verify error handling
# Check Redis cache was set to error status
from extensions.ext_redis import redis_client
cache_key = f"segment_batch_import_{job_id}"
cache_value = redis_client.get(cache_key)
assert cache_value == b"error"
# Verify no segments were created
from extensions.ext_database import db
segments = db.session.query(DocumentSegment).all()
assert len(segments) == 0
# Verify dataset remains unchanged (no segments were added to the dataset)
db.session.refresh(dataset)
segments_for_dataset = db.session.query(DocumentSegment).filter_by(dataset_id=dataset.id).all()
assert len(segments_for_dataset) == 0
def test_batch_create_segment_to_index_task_document_not_available(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test task failure when document is not available for indexing.
This test verifies that the task properly handles error cases:
1. Fails when document is disabled
2. Fails when document is archived
3. Fails when document indexing status is not completed
4. Sets appropriate Redis cache status
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
# Create document with various unavailable states
test_cases = [
# Disabled document
Document(
tenant_id=tenant.id,
dataset_id=dataset.id,
position=1,
data_source_type="upload_file",
batch="test_batch",
name="disabled_document",
created_from="upload_file",
created_by=account.id,
indexing_status="completed",
enabled=False, # Document is disabled
archived=False,
doc_form="text_model",
word_count=0,
),
# Archived document
Document(
tenant_id=tenant.id,
dataset_id=dataset.id,
position=2,
data_source_type="upload_file",
batch="test_batch",
name="archived_document",
created_from="upload_file",
created_by=account.id,
indexing_status="completed",
enabled=True,
archived=True, # Document is archived
doc_form="text_model",
word_count=0,
),
# Document with incomplete indexing
Document(
tenant_id=tenant.id,
dataset_id=dataset.id,
position=3,
data_source_type="upload_file",
batch="test_batch",
name="incomplete_document",
created_from="upload_file",
created_by=account.id,
indexing_status="indexing", # Not completed
enabled=True,
archived=False,
doc_form="text_model",
word_count=0,
),
]
from extensions.ext_database import db
for document in test_cases:
db.session.add(document)
db.session.commit()
# Test each unavailable document
for document in test_cases:
job_id = str(uuid.uuid4())
batch_create_segment_to_index_task(
job_id=job_id,
upload_file_id=upload_file.id,
dataset_id=dataset.id,
document_id=document.id,
tenant_id=tenant.id,
user_id=account.id,
)
# Verify error handling for each case
from extensions.ext_redis import redis_client
cache_key = f"segment_batch_import_{job_id}"
cache_value = redis_client.get(cache_key)
assert cache_value == b"error"
# Verify no segments were created
segments = db.session.query(DocumentSegment).filter_by(document_id=document.id).all()
assert len(segments) == 0
def test_batch_create_segment_to_index_task_upload_file_not_found(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test task failure when upload file does not exist.
This test verifies that the task properly handles error cases:
1. Fails gracefully when upload file is not found
2. Sets appropriate Redis cache status
3. Maintains database integrity
4. Logs appropriate error information
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
# Use non-existent upload file ID
non_existent_upload_file_id = str(uuid.uuid4())
# Execute the task with non-existent upload file
job_id = str(uuid.uuid4())
batch_create_segment_to_index_task(
job_id=job_id,
upload_file_id=non_existent_upload_file_id,
dataset_id=dataset.id,
document_id=document.id,
tenant_id=tenant.id,
user_id=account.id,
)
# Verify error handling
# Check Redis cache was set to error status
from extensions.ext_redis import redis_client
cache_key = f"segment_batch_import_{job_id}"
cache_value = redis_client.get(cache_key)
assert cache_value == b"error"
# Verify no segments were created
from extensions.ext_database import db
segments = db.session.query(DocumentSegment).all()
assert len(segments) == 0
# Verify document remains unchanged
db.session.refresh(document)
assert document.word_count == 0
def test_batch_create_segment_to_index_task_empty_csv_file(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test task failure when CSV file is empty.
This test verifies that the task properly handles error cases:
1. Fails when CSV file contains no data
2. Sets appropriate Redis cache status
3. Maintains database integrity
4. Logs appropriate error information
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
# Create empty CSV content
empty_csv_content = "content\n" # Only header, no data rows
# Mock storage to return empty CSV content
mock_storage = mock_external_service_dependencies["storage"]
def mock_download(key, file_path):
Path(file_path).write_text(empty_csv_content, encoding="utf-8")
mock_storage.download.side_effect = mock_download
# Execute the task
job_id = str(uuid.uuid4())
batch_create_segment_to_index_task(
job_id=job_id,
upload_file_id=upload_file.id,
dataset_id=dataset.id,
document_id=document.id,
tenant_id=tenant.id,
user_id=account.id,
)
# Verify error handling
# Check Redis cache was set to error status
from extensions.ext_redis import redis_client
cache_key = f"segment_batch_import_{job_id}"
cache_value = redis_client.get(cache_key)
assert cache_value == b"error"
# Verify no segments were created
from extensions.ext_database import db
segments = db.session.query(DocumentSegment).all()
assert len(segments) == 0
# Verify document remains unchanged
db.session.refresh(document)
assert document.word_count == 0
def test_batch_create_segment_to_index_task_position_calculation(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test proper position calculation for segments when existing segments exist.
This test verifies that the task correctly:
1. Calculates positions for new segments based on existing ones
2. Handles position increment logic properly
3. Maintains proper segment ordering
4. Works with existing segment data
"""
# Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(db_session_with_containers, account, tenant)
document = self._create_test_document(db_session_with_containers, account, tenant, dataset)
upload_file = self._create_test_upload_file(db_session_with_containers, account, tenant)
# Create existing segments to test position calculation
existing_segments = []
for i in range(3):
segment = DocumentSegment(
tenant_id=tenant.id,
dataset_id=dataset.id,
document_id=document.id,
position=i + 1,
content=f"Existing segment {i + 1}",
word_count=len(f"Existing segment {i + 1}"),
tokens=10,
created_by=account.id,
status="completed",
index_node_id=str(uuid.uuid4()),
index_node_hash=f"hash_{i}",
)
existing_segments.append(segment)
from extensions.ext_database import db
for segment in existing_segments:
db.session.add(segment)
db.session.commit()
# Create CSV content
csv_content = self._create_test_csv_content("text_model")
# Mock storage to return our CSV content
mock_storage = mock_external_service_dependencies["storage"]
def mock_download(key, file_path):
Path(file_path).write_text(csv_content, encoding="utf-8")
mock_storage.download.side_effect = mock_download
# Execute the task
job_id = str(uuid.uuid4())
batch_create_segment_to_index_task(
job_id=job_id,
upload_file_id=upload_file.id,
dataset_id=dataset.id,
document_id=document.id,
tenant_id=tenant.id,
user_id=account.id,
)
# Verify results
# Check that new segments were created with correct positions
all_segments = (
db.session.query(DocumentSegment)
.filter_by(document_id=document.id)
.order_by(DocumentSegment.position)
.all()
)
assert len(all_segments) == 6 # 3 existing + 3 new
# Verify position ordering
for i, segment in enumerate(all_segments):
assert segment.position == i + 1
# Verify new segments have correct positions (4, 5, 6)
new_segments = all_segments[3:]
for i, segment in enumerate(new_segments):
expected_position = 4 + i # Should start at position 4
assert segment.position == expected_position
assert segment.status == "completed"
assert segment.indexing_at is not None
assert segment.completed_at is not None
# Check that document word count was updated
db.session.refresh(document)
assert document.word_count > 0
# Verify vector service was called
mock_vector_service = mock_external_service_dependencies["vector_service"]
mock_vector_service.create_segments_vector.assert_called_once()
# Check Redis cache was set
from extensions.ext_redis import redis_client
cache_key = f"segment_batch_import_{job_id}"
cache_value = redis_client.get(cache_key)
assert cache_value == b"completed"

1144
api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py Normal file
View File

File diff suppressed because it is too large Load Diff

1153
api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py Normal file
View File

File diff suppressed because it is too large Load Diff

1099
api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py Normal file
View File

File diff suppressed because it is too large Load Diff

1391
api/tests/test_containers_integration_tests/tasks/test_deal_dataset_vector_index_task.py Normal file
View File

File diff suppressed because it is too large Load Diff

583
api/tests/test_containers_integration_tests/tasks/test_delete_segment_from_index_task.py Normal file
View File

@ -0,0 +1,583 @@
"""
TestContainers-based integration tests for delete_segment_from_index_task.
This module provides comprehensive integration testing for the delete_segment_from_index_task
using TestContainers to ensure realistic database interactions and proper isolation.
The task is responsible for removing document segments from the vector index when segments
are deleted from the dataset.
"""
import logging
from unittest.mock import MagicMock, patch
from faker import Faker
from core.rag.index_processor.constant.index_type import IndexType
from models import Account, Dataset, Document, DocumentSegment, Tenant
from tasks.delete_segment_from_index_task import delete_segment_from_index_task
logger = logging.getLogger(__name__)
class TestDeleteSegmentFromIndexTask:
"""
Comprehensive integration tests for delete_segment_from_index_task using testcontainers.
This test class covers all major functionality of the delete_segment_from_index_task:
- Successful segment deletion from index
- Dataset not found scenarios
- Document not found scenarios
- Document status validation (disabled, archived, not completed)
- Index processor integration and cleanup
- Exception handling and error scenarios
- Performance and timing verification
All tests use the testcontainers infrastructure to ensure proper database isolation
and realistic testing environment with actual database interactions.
"""
def _create_test_tenant(self, db_session_with_containers, fake=None):
"""
Helper method to create a test tenant with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
fake: Faker instance for generating test data
Returns:
Tenant: Created test tenant instance
"""
fake = fake or Faker()
tenant = Tenant()
tenant.id = fake.uuid4()
tenant.name = f"Test Tenant {fake.company()}"
tenant.plan = "basic"
tenant.status = "active"
tenant.created_at = fake.date_time_this_year()
tenant.updated_at = tenant.created_at
db_session_with_containers.add(tenant)
db_session_with_containers.commit()
return tenant
def _create_test_account(self, db_session_with_containers, tenant, fake=None):
"""
Helper method to create a test account with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
tenant: Tenant instance for the account
fake: Faker instance for generating test data
Returns:
Account: Created test account instance
"""
fake = fake or Faker()
account = Account()
account.id = fake.uuid4()
account.email = fake.email()
account.name = fake.name()
account.avatar_url = fake.url()
account.tenant_id = tenant.id
account.status = "active"
account.type = "normal"
account.role = "owner"
account.interface_language = "en-US"
account.created_at = fake.date_time_this_year()
account.updated_at = account.created_at
db_session_with_containers.add(account)
db_session_with_containers.commit()
return account
def _create_test_dataset(self, db_session_with_containers, tenant, account, fake=None):
"""
Helper method to create a test dataset with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
tenant: Tenant instance for the dataset
account: Account instance for the dataset
fake: Faker instance for generating test data
Returns:
Dataset: Created test dataset instance
"""
fake = fake or Faker()
dataset = Dataset()
dataset.id = fake.uuid4()
dataset.tenant_id = tenant.id
dataset.name = f"Test Dataset {fake.word()}"
dataset.description = fake.text(max_nb_chars=200)
dataset.provider = "vendor"
dataset.permission = "only_me"
dataset.data_source_type = "upload_file"
dataset.indexing_technique = "high_quality"
dataset.index_struct = '{"type": "paragraph"}'
dataset.created_by = account.id
dataset.created_at = fake.date_time_this_year()
dataset.updated_by = account.id
dataset.updated_at = dataset.created_at
dataset.embedding_model = "text-embedding-ada-002"
dataset.embedding_model_provider = "openai"
dataset.built_in_field_enabled = False
db_session_with_containers.add(dataset)
db_session_with_containers.commit()
return dataset
def _create_test_document(self, db_session_with_containers, dataset, account, fake=None, **kwargs):
"""
Helper method to create a test document with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
dataset: Dataset instance for the document
account: Account instance for the document
fake: Faker instance for generating test data
**kwargs: Additional document attributes to override defaults
Returns:
Document: Created test document instance
"""
fake = fake or Faker()
document = Document()
document.id = fake.uuid4()
document.tenant_id = dataset.tenant_id
document.dataset_id = dataset.id
document.position = kwargs.get("position", 1)
document.data_source_type = kwargs.get("data_source_type", "upload_file")
document.data_source_info = kwargs.get("data_source_info", "{}")
document.batch = kwargs.get("batch", fake.uuid4())
document.name = kwargs.get("name", f"Test Document {fake.word()}")
document.created_from = kwargs.get("created_from", "api")
document.created_by = account.id
document.created_at = fake.date_time_this_year()
document.processing_started_at = kwargs.get("processing_started_at", fake.date_time_this_year())
document.file_id = kwargs.get("file_id", fake.uuid4())
document.word_count = kwargs.get("word_count", fake.random_int(min=100, max=1000))
document.parsing_completed_at = kwargs.get("parsing_completed_at", fake.date_time_this_year())
document.cleaning_completed_at = kwargs.get("cleaning_completed_at", fake.date_time_this_year())
document.splitting_completed_at = kwargs.get("splitting_completed_at", fake.date_time_this_year())
document.tokens = kwargs.get("tokens", fake.random_int(min=50, max=500))
document.indexing_latency = kwargs.get("indexing_latency", fake.random_number(digits=3))
document.completed_at = kwargs.get("completed_at", fake.date_time_this_year())
document.is_paused = kwargs.get("is_paused", False)
document.indexing_status = kwargs.get("indexing_status", "completed")
document.enabled = kwargs.get("enabled", True)
document.archived = kwargs.get("archived", False)
document.updated_at = fake.date_time_this_year()
document.doc_type = kwargs.get("doc_type", "text")
document.doc_metadata = kwargs.get("doc_metadata", {})
document.doc_form = kwargs.get("doc_form", IndexType.PARAGRAPH_INDEX)
document.doc_language = kwargs.get("doc_language", "en")
db_session_with_containers.add(document)
db_session_with_containers.commit()
return document
def _create_test_document_segments(self, db_session_with_containers, document, account, count=3, fake=None):
"""
Helper method to create test document segments with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
document: Document instance for the segments
account: Account instance for the segments
count: Number of segments to create
fake: Faker instance for generating test data
Returns:
list[DocumentSegment]: List of created test document segment instances
"""
fake = fake or Faker()
segments = []
for i in range(count):
segment = DocumentSegment()
segment.id = fake.uuid4()
segment.tenant_id = document.tenant_id
segment.dataset_id = document.dataset_id
segment.document_id = document.id
segment.position = i + 1
segment.content = f"Test segment content {i + 1}: {fake.text(max_nb_chars=200)}"
segment.answer = f"Test segment answer {i + 1}: {fake.text(max_nb_chars=100)}"
segment.word_count = fake.random_int(min=10, max=100)
segment.tokens = fake.random_int(min=5, max=50)
segment.keywords = [fake.word() for _ in range(3)]
segment.index_node_id = f"node_{fake.uuid4()}"
segment.index_node_hash = fake.sha256()
segment.hit_count = 0
segment.enabled = True
segment.status = "completed"
segment.created_by = account.id
segment.created_at = fake.date_time_this_year()
segment.updated_by = account.id
segment.updated_at = segment.created_at
db_session_with_containers.add(segment)
segments.append(segment)
db_session_with_containers.commit()
return segments
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
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.
This test verifies:
- Proper task execution with valid dataset and document
- Index processor factory initialization with correct document form
- Index processor clean method called with correct parameters
- Database session properly closed after execution
- Task completes without exceptions
"""
fake = Faker()
# Create test data
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_document_segments(db_session_with_containers, document, account, 3, fake)
# Extract index node IDs for the task
index_node_ids = [segment.index_node_id for segment in segments]
# Mock the index processor
mock_processor = MagicMock()
mock_index_processor_factory.return_value.init_index_processor.return_value = mock_processor
# Execute the task
result = delete_segment_from_index_task(index_node_ids, dataset.id, document.id)
# Verify the task completed successfully
assert result is None # Task should return None on success
# Verify index processor factory was called with correct document form
mock_index_processor_factory.assert_called_once_with(document.doc_form)
# Verify index processor clean method was called with correct parameters
# Note: We can't directly compare Dataset objects as they are different instances
# from database queries, so we verify the call was made and check the parameters
assert mock_processor.clean.call_count == 1
call_args = mock_processor.clean.call_args
assert call_args[0][0].id == dataset.id # Verify dataset ID matches
assert call_args[0][1] == index_node_ids # Verify index node IDs match
assert call_args[1]["with_keywords"] is True
assert call_args[1]["delete_child_chunks"] is True
def test_delete_segment_from_index_task_dataset_not_found(self, db_session_with_containers):
"""
Test task behavior when dataset is not found.
This test verifies:
- Task handles missing dataset gracefully
- No index processor operations are attempted
- Task returns early without exceptions
- Database session is properly closed
"""
fake = Faker()
non_existent_dataset_id = fake.uuid4()
non_existent_document_id = fake.uuid4()
index_node_ids = [f"node_{fake.uuid4()}" for _ in range(3)]
# Execute the task with non-existent dataset
result = delete_segment_from_index_task(index_node_ids, non_existent_dataset_id, non_existent_document_id)
# Verify the task completed without exceptions
assert result is None # Task should return None when dataset not found
def test_delete_segment_from_index_task_document_not_found(self, db_session_with_containers):
"""
Test task behavior when document is not found.
This test verifies:
- Task handles missing document gracefully
- No index processor operations are attempted
- Task returns early without exceptions
- Database session is properly closed
"""
fake = Faker()
# Create test data
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
non_existent_document_id = fake.uuid4()
index_node_ids = [f"node_{fake.uuid4()}" for _ in range(3)]
# Execute the task with non-existent document
result = delete_segment_from_index_task(index_node_ids, dataset.id, non_existent_document_id)
# Verify the task completed without exceptions
assert result is None # Task should return None when document not found
def test_delete_segment_from_index_task_document_disabled(self, db_session_with_containers):
"""
Test task behavior when document is disabled.
This test verifies:
- Task handles disabled document gracefully
- No index processor operations are attempted
- Task returns early without exceptions
- Database session is properly closed
"""
fake = Faker()
# Create test data with disabled document
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake, enabled=False)
segments = self._create_test_document_segments(db_session_with_containers, document, account, 3, fake)
index_node_ids = [segment.index_node_id for segment in segments]
# Execute the task with disabled document
result = delete_segment_from_index_task(index_node_ids, dataset.id, document.id)
# Verify the task completed without exceptions
assert result is None # Task should return None when document is disabled
def test_delete_segment_from_index_task_document_archived(self, db_session_with_containers):
"""
Test task behavior when document is archived.
This test verifies:
- Task handles archived document gracefully
- No index processor operations are attempted
- Task returns early without exceptions
- Database session is properly closed
"""
fake = Faker()
# Create test data with archived document
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake, archived=True)
segments = self._create_test_document_segments(db_session_with_containers, document, account, 3, fake)
index_node_ids = [segment.index_node_id for segment in segments]
# Execute the task with archived document
result = delete_segment_from_index_task(index_node_ids, dataset.id, document.id)
# Verify the task completed without exceptions
assert result is None # Task should return None when document is archived
def test_delete_segment_from_index_task_document_not_completed(self, db_session_with_containers):
"""
Test task behavior when document indexing is not completed.
This test verifies:
- Task handles incomplete indexing status gracefully
- No index processor operations are attempted
- Task returns early without exceptions
- Database session is properly closed
"""
fake = Faker()
# Create test data with incomplete indexing
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
document = self._create_test_document(
db_session_with_containers, dataset, account, fake, indexing_status="indexing"
)
segments = self._create_test_document_segments(db_session_with_containers, document, account, 3, fake)
index_node_ids = [segment.index_node_id for segment in segments]
# Execute the task with incomplete indexing
result = delete_segment_from_index_task(index_node_ids, dataset.id, document.id)
# 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")
def test_delete_segment_from_index_task_index_processor_clean(
self, mock_index_processor_factory, db_session_with_containers
):
"""
Test index processor clean method integration with different document forms.
This test verifies:
- Index processor factory creates correct processor for different document forms
- Clean method is called with proper parameters for each document form
- Task handles different index types correctly
- Database session is properly managed
"""
fake = Faker()
# Test different document forms
document_forms = [IndexType.PARAGRAPH_INDEX, IndexType.QA_INDEX, IndexType.PARENT_CHILD_INDEX]
for doc_form in document_forms:
# Create test data for each document form
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake, doc_form=doc_form)
segments = self._create_test_document_segments(db_session_with_containers, document, account, 2, fake)
index_node_ids = [segment.index_node_id for segment in segments]
# Mock the index processor
mock_processor = MagicMock()
mock_index_processor_factory.return_value.init_index_processor.return_value = mock_processor
# Execute the task
result = delete_segment_from_index_task(index_node_ids, dataset.id, document.id)
# Verify the task completed successfully
assert result is None
# Verify index processor factory was called with correct document form
mock_index_processor_factory.assert_called_with(doc_form)
# Verify index processor clean method was called with correct parameters
assert mock_processor.clean.call_count == 1
call_args = mock_processor.clean.call_args
assert call_args[0][0].id == dataset.id # Verify dataset ID matches
assert call_args[0][1] == index_node_ids # Verify index node IDs match
assert call_args[1]["with_keywords"] is True
assert call_args[1]["delete_child_chunks"] is True
# Reset mocks for next iteration
mock_index_processor_factory.reset_mock()
mock_processor.reset_mock()
@patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
def test_delete_segment_from_index_task_exception_handling(
self, mock_index_processor_factory, db_session_with_containers
):
"""
Test exception handling in the task.
This test verifies:
- Task handles index processor exceptions gracefully
- Database session is properly closed even when exceptions occur
- Task logs exceptions appropriately
- No unhandled exceptions are raised
"""
fake = Faker()
# Create test data
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_document_segments(db_session_with_containers, document, account, 3, fake)
index_node_ids = [segment.index_node_id for segment in segments]
# Mock the index processor to raise an exception
mock_processor = MagicMock()
mock_processor.clean.side_effect = Exception("Index processor error")
mock_index_processor_factory.return_value.init_index_processor.return_value = mock_processor
# Execute the task - should not raise exception
result = delete_segment_from_index_task(index_node_ids, dataset.id, document.id)
# Verify the task completed without raising exceptions
assert result is None # Task should return None even when exceptions occur
# Verify index processor clean method was called
assert mock_processor.clean.call_count == 1
call_args = mock_processor.clean.call_args
assert call_args[0][0].id == dataset.id # Verify dataset ID matches
assert call_args[0][1] == index_node_ids # Verify index node IDs match
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")
def test_delete_segment_from_index_task_empty_index_node_ids(
self, mock_index_processor_factory, db_session_with_containers
):
"""
Test task behavior with empty index node IDs list.
This test verifies:
- Task handles empty index node IDs gracefully
- Index processor clean method is called with empty list
- Task completes successfully
- Database session is properly managed
"""
fake = Faker()
# Create test data
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
# Use empty index node IDs
index_node_ids = []
# Mock the index processor
mock_processor = MagicMock()
mock_index_processor_factory.return_value.init_index_processor.return_value = mock_processor
# Execute the task
result = delete_segment_from_index_task(index_node_ids, dataset.id, document.id)
# Verify the task completed successfully
assert result is None
# Verify index processor clean method was called with empty list
assert mock_processor.clean.call_count == 1
call_args = mock_processor.clean.call_args
assert call_args[0][0].id == dataset.id # Verify dataset ID matches
assert call_args[0][1] == index_node_ids # Verify index node IDs match (empty list)
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")
def test_delete_segment_from_index_task_large_index_node_ids(
self, mock_index_processor_factory, db_session_with_containers
):
"""
Test task behavior with large number of index node IDs.
This test verifies:
- Task handles large lists of index node IDs efficiently
- Index processor clean method is called with all node IDs
- Task completes successfully with large datasets
- Database session is properly managed
"""
fake = Faker()
# Create test data
tenant = self._create_test_tenant(db_session_with_containers, fake)
account = self._create_test_account(db_session_with_containers, tenant, fake)
dataset = self._create_test_dataset(db_session_with_containers, tenant, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
# Create large number of segments
segments = self._create_test_document_segments(db_session_with_containers, document, account, 50, fake)
index_node_ids = [segment.index_node_id for segment in segments]
# Mock the index processor
mock_processor = MagicMock()
mock_index_processor_factory.return_value.init_index_processor.return_value = mock_processor
# Execute the task
result = delete_segment_from_index_task(index_node_ids, dataset.id, document.id)
# Verify the task completed successfully
assert result is None
# Verify index processor clean method was called with all node IDs
assert mock_processor.clean.call_count == 1
call_args = mock_processor.clean.call_args
assert call_args[0][0].id == dataset.id # Verify dataset ID matches
assert call_args[0][1] == index_node_ids # Verify index node IDs match
assert call_args[1]["with_keywords"] is True
assert call_args[1]["delete_child_chunks"] is True
# Verify all node IDs were passed
assert len(call_args[0][1]) == 50

615
api/tests/test_containers_integration_tests/tasks/test_disable_segment_from_index_task.py Normal file
View File

@ -0,0 +1,615 @@
"""
Integration tests for disable_segment_from_index_task using TestContainers.
This module provides comprehensive integration tests for the disable_segment_from_index_task
using real database and Redis containers to ensure the task works correctly with actual
data and external dependencies.
"""
import logging
import time
from datetime import UTC, datetime
from unittest.mock import patch
import pytest
from faker import Faker
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, Document, DocumentSegment
from tasks.disable_segment_from_index_task import disable_segment_from_index_task
logger = logging.getLogger(__name__)
class TestDisableSegmentFromIndexTask:
"""Integration tests for disable_segment_from_index_task using testcontainers."""
@pytest.fixture
def mock_index_processor(self):
"""Mock IndexProcessorFactory and its clean method."""
with patch("tasks.disable_segment_from_index_task.IndexProcessorFactory") as mock_factory:
mock_processor = mock_factory.return_value.init_index_processor.return_value
mock_processor.clean.return_value = None
yield mock_processor
def _create_test_account_and_tenant(self, db_session_with_containers) -> tuple[Account, Tenant]:
"""
Helper method to create a test account and tenant for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
Returns:
tuple: (account, tenant) - Created account and tenant instances
"""
fake = Faker()
# Create account
account = Account(
email=fake.email(),
name=fake.name(),
interface_language="en-US",
status="active",
)
db.session.add(account)
db.session.commit()
# Create tenant
tenant = Tenant(
name=fake.company(),
status="normal",
plan="basic",
)
db.session.add(tenant)
db.session.commit()
# Create tenant-account join with owner role
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER.value,
current=True,
)
db.session.add(join)
db.session.commit()
# Set current tenant for account
account.current_tenant = tenant
return account, tenant
def _create_test_dataset(self, tenant: Tenant, account: Account) -> Dataset:
"""
Helper method to create a test dataset.
Args:
tenant: Tenant instance
account: Account instance
Returns:
Dataset: Created dataset instance
"""
fake = Faker()
dataset = Dataset(
tenant_id=tenant.id,
name=fake.sentence(nb_words=3),
description=fake.text(max_nb_chars=200),
data_source_type="upload_file",
indexing_technique="high_quality",
created_by=account.id,
)
db.session.add(dataset)
db.session.commit()
return dataset
def _create_test_document(
self, dataset: Dataset, tenant: Tenant, account: Account, doc_form: str = "text_model"
) -> Document:
"""
Helper method to create a test document.
Args:
dataset: Dataset instance
tenant: Tenant instance
account: Account instance
doc_form: Document form type
Returns:
Document: Created document instance
"""
fake = Faker()
document = Document(
tenant_id=tenant.id,
dataset_id=dataset.id,
position=1,
data_source_type="upload_file",
batch=fake.uuid4(),
name=fake.file_name(),
created_from="api",
created_by=account.id,
indexing_status="completed",
enabled=True,
archived=False,
doc_form=doc_form,
word_count=1000,
tokens=500,
completed_at=datetime.now(UTC),
)
db.session.add(document)
db.session.commit()
return document
def _create_test_segment(
self,
document: Document,
dataset: Dataset,
tenant: Tenant,
account: Account,
status: str = "completed",
enabled: bool = True,
) -> DocumentSegment:
"""
Helper method to create a test document segment.
Args:
document: Document instance
dataset: Dataset instance
tenant: Tenant instance
account: Account instance
status: Segment status
enabled: Whether segment is enabled
Returns:
DocumentSegment: Created segment instance
"""
fake = Faker()
segment = DocumentSegment(
tenant_id=tenant.id,
dataset_id=dataset.id,
document_id=document.id,
position=1,
content=fake.text(max_nb_chars=500),
word_count=100,
tokens=50,
index_node_id=fake.uuid4(),
index_node_hash=fake.sha256(),
status=status,
enabled=enabled,
created_by=account.id,
completed_at=datetime.now(UTC) if status == "completed" else None,
)
db.session.add(segment)
db.session.commit()
return segment
def test_disable_segment_success(self, db_session_with_containers, mock_index_processor):
"""
Test successful segment disabling from index.
This test verifies:
- Segment is found and validated
- Index processor clean method is called with correct parameters
- Redis cache is cleared
- Task completes successfully
"""
# Arrange: Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segment = self._create_test_segment(document, dataset, tenant, account)
# Set up Redis cache
indexing_cache_key = f"segment_{segment.id}_indexing"
redis_client.setex(indexing_cache_key, 600, 1)
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task completed successfully
assert result is None # Task returns None on success
# Verify index processor was called correctly
mock_index_processor.clean.assert_called_once()
call_args = mock_index_processor.clean.call_args
assert call_args[0][0].id == dataset.id # Check dataset ID
assert call_args[0][1] == [segment.index_node_id] # Check index node IDs
# Verify Redis cache was cleared
assert redis_client.get(indexing_cache_key) is None
# Verify segment is still in database
db.session.refresh(segment)
assert segment.id is not None
def test_disable_segment_not_found(self, db_session_with_containers, mock_index_processor):
"""
Test handling when segment is not found.
This test verifies:
- Task handles non-existent segment gracefully
- No index processor operations are performed
- Task returns early without errors
"""
# Arrange: Use a non-existent segment ID
fake = Faker()
non_existent_segment_id = fake.uuid4()
# Act: Execute the task with non-existent segment
result = disable_segment_from_index_task(non_existent_segment_id)
# Assert: Verify the task handled the error gracefully
assert result is None
# Verify index processor was not called
mock_index_processor.clean.assert_not_called()
def test_disable_segment_not_completed(self, db_session_with_containers, mock_index_processor):
"""
Test handling when segment is not in completed status.
This test verifies:
- Task rejects segments that are not completed
- No index processor operations are performed
- Task returns early without errors
"""
# Arrange: Create test data with non-completed segment
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segment = self._create_test_segment(document, dataset, tenant, account, status="indexing", enabled=True)
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task handled the invalid status gracefully
assert result is None
# Verify index processor was not called
mock_index_processor.clean.assert_not_called()
def test_disable_segment_no_dataset(self, db_session_with_containers, mock_index_processor):
"""
Test handling when segment has no associated dataset.
This test verifies:
- Task handles segments without dataset gracefully
- No index processor operations are performed
- Task returns early without errors
"""
# Arrange: Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segment = self._create_test_segment(document, dataset, tenant, account)
# Manually remove dataset association
segment.dataset_id = "00000000-0000-0000-0000-000000000000"
db.session.commit()
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task handled the missing dataset gracefully
assert result is None
# Verify index processor was not called
mock_index_processor.clean.assert_not_called()
def test_disable_segment_no_document(self, db_session_with_containers, mock_index_processor):
"""
Test handling when segment has no associated document.
This test verifies:
- Task handles segments without document gracefully
- No index processor operations are performed
- Task returns early without errors
"""
# Arrange: Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segment = self._create_test_segment(document, dataset, tenant, account)
# Manually remove document association
segment.document_id = "00000000-0000-0000-0000-000000000000"
db.session.commit()
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task handled the missing document gracefully
assert result is None
# Verify index processor was not called
mock_index_processor.clean.assert_not_called()
def test_disable_segment_document_disabled(self, db_session_with_containers, mock_index_processor):
"""
Test handling when document is disabled.
This test verifies:
- Task handles disabled documents gracefully
- No index processor operations are performed
- Task returns early without errors
"""
# Arrange: Create test data with disabled document
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
document.enabled = False
db.session.commit()
segment = self._create_test_segment(document, dataset, tenant, account)
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task handled the disabled document gracefully
assert result is None
# Verify index processor was not called
mock_index_processor.clean.assert_not_called()
def test_disable_segment_document_archived(self, db_session_with_containers, mock_index_processor):
"""
Test handling when document is archived.
This test verifies:
- Task handles archived documents gracefully
- No index processor operations are performed
- Task returns early without errors
"""
# Arrange: Create test data with archived document
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
document.archived = True
db.session.commit()
segment = self._create_test_segment(document, dataset, tenant, account)
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task handled the archived document gracefully
assert result is None
# Verify index processor was not called
mock_index_processor.clean.assert_not_called()
def test_disable_segment_document_indexing_not_completed(self, db_session_with_containers, mock_index_processor):
"""
Test handling when document indexing is not completed.
This test verifies:
- Task handles documents with incomplete indexing gracefully
- No index processor operations are performed
- Task returns early without errors
"""
# Arrange: Create test data with incomplete indexing
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
document.indexing_status = "indexing"
db.session.commit()
segment = self._create_test_segment(document, dataset, tenant, account)
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task handled the incomplete indexing gracefully
assert result is None
# Verify index processor was not called
mock_index_processor.clean.assert_not_called()
def test_disable_segment_index_processor_exception(self, db_session_with_containers, mock_index_processor):
"""
Test handling when index processor raises an exception.
This test verifies:
- Task handles index processor exceptions gracefully
- Segment is re-enabled on failure
- Redis cache is still cleared
- Database changes are committed
"""
# Arrange: Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segment = self._create_test_segment(document, dataset, tenant, account)
# Set up Redis cache
indexing_cache_key = f"segment_{segment.id}_indexing"
redis_client.setex(indexing_cache_key, 600, 1)
# Configure mock to raise exception
mock_index_processor.clean.side_effect = Exception("Index processor error")
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task handled the exception gracefully
assert result is None
# Verify index processor was called
mock_index_processor.clean.assert_called_once()
call_args = mock_index_processor.clean.call_args
# Check that the call was made with the correct parameters
assert len(call_args[0]) == 2 # Check two arguments were passed
assert call_args[0][1] == [segment.index_node_id] # Check index node IDs
# Verify segment was re-enabled
db.session.refresh(segment)
assert segment.enabled is True
# Verify Redis cache was still cleared
assert redis_client.get(indexing_cache_key) is None
def test_disable_segment_different_doc_forms(self, db_session_with_containers, mock_index_processor):
"""
Test disabling segments with different document forms.
This test verifies:
- Task works with different document form types
- Correct index processor is initialized for each form
- Index processor clean method is called correctly
"""
# Test different document forms
doc_forms = ["text_model", "qa_model", "table_model"]
for doc_form in doc_forms:
# Arrange: Create test data for each form
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account, doc_form=doc_form)
segment = self._create_test_segment(document, dataset, tenant, account)
# Reset mock for each iteration
mock_index_processor.reset_mock()
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify the task completed successfully
assert result is None
# Verify correct index processor was initialized
mock_index_processor.clean.assert_called_once()
call_args = mock_index_processor.clean.call_args
assert call_args[0][0].id == dataset.id # Check dataset ID
assert call_args[0][1] == [segment.index_node_id] # Check index node IDs
def test_disable_segment_redis_cache_handling(self, db_session_with_containers, mock_index_processor):
"""
Test Redis cache handling during segment disabling.
This test verifies:
- Redis cache is properly set before task execution
- Cache is cleared after task completion
- Cache handling works with different scenarios
"""
# Arrange: Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segment = self._create_test_segment(document, dataset, tenant, account)
# Test with cache present
indexing_cache_key = f"segment_{segment.id}_indexing"
redis_client.setex(indexing_cache_key, 600, 1)
assert redis_client.get(indexing_cache_key) is not None
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify cache was cleared
assert result is None
assert redis_client.get(indexing_cache_key) is None
# Test with no cache present
segment2 = self._create_test_segment(document, dataset, tenant, account)
result2 = disable_segment_from_index_task(segment2.id)
# Assert: Verify task still works without cache
assert result2 is None
def test_disable_segment_performance_timing(self, db_session_with_containers, mock_index_processor):
"""
Test performance timing of segment disabling task.
This test verifies:
- Task execution time is reasonable
- Performance logging works correctly
- Task completes within expected time bounds
"""
# Arrange: Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segment = self._create_test_segment(document, dataset, tenant, account)
# Act: Execute the task and measure time
start_time = time.perf_counter()
result = disable_segment_from_index_task(segment.id)
end_time = time.perf_counter()
# Assert: Verify task completed successfully and timing is reasonable
assert result is None
execution_time = end_time - start_time
assert execution_time < 5.0 # Should complete within 5 seconds
def test_disable_segment_database_session_management(self, db_session_with_containers, mock_index_processor):
"""
Test database session management during task execution.
This test verifies:
- Database sessions are properly managed
- Sessions are closed after task completion
- No session leaks occur
"""
# Arrange: Create test data
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segment = self._create_test_segment(document, dataset, tenant, account)
# Act: Execute the task
result = disable_segment_from_index_task(segment.id)
# Assert: Verify task completed and session management worked
assert result is None
# Verify segment is still accessible (session was properly managed)
db.session.refresh(segment)
assert segment.id is not None
def test_disable_segment_concurrent_execution(self, db_session_with_containers, mock_index_processor):
"""
Test concurrent execution of segment disabling tasks.
This test verifies:
- Multiple tasks can run concurrently
- Each task processes its own segment correctly
- No interference between concurrent tasks
"""
# Arrange: Create multiple test segments
account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
dataset = self._create_test_dataset(tenant, account)
document = self._create_test_document(dataset, tenant, account)
segments = []
for i in range(3):
segment = self._create_test_segment(document, dataset, tenant, account)
segments.append(segment)
# Act: Execute tasks concurrently (simulated)
results = []
for segment in segments:
result = disable_segment_from_index_task(segment.id)
results.append(result)
# Assert: Verify all tasks completed successfully
assert all(result is None for result in results)
# Verify all segments were processed
assert mock_index_processor.clean.call_count == len(segments)
# Verify each segment was processed with correct parameters
for segment in segments:
# Check that clean was called with this segment's dataset and index_node_id
found = False
for call in mock_index_processor.clean.call_args_list:
if call[0][0].id == dataset.id and call[0][1] == [segment.index_node_id]:
found = True
break
assert found, f"Segment {segment.id} was not processed correctly"

729
api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py Normal file
View File

@ -0,0 +1,729 @@
"""
TestContainers-based integration tests for disable_segments_from_index_task.
This module provides comprehensive integration testing for the disable_segments_from_index_task
using TestContainers to ensure realistic database interactions and proper isolation.
The task is responsible for removing document segments from the search index when they are disabled.
"""
from unittest.mock import MagicMock, patch
from faker import Faker
from models import Account, Dataset, DocumentSegment
from models import Document as DatasetDocument
from models.dataset import DatasetProcessRule
from tasks.disable_segments_from_index_task import disable_segments_from_index_task
class TestDisableSegmentsFromIndexTask:
"""
Comprehensive integration tests for disable_segments_from_index_task using testcontainers.
This test class covers all major functionality of the disable_segments_from_index_task:
- Successful segment disabling with proper index cleanup
- Error handling for various edge cases
- Database state validation after task execution
- Redis cache cleanup verification
- Index processor integration testing
All tests use the testcontainers infrastructure to ensure proper database isolation
and realistic testing environment with actual database interactions.
"""
def _create_test_account(self, db_session_with_containers, fake=None):
"""
Helper method to create a test account with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
fake: Faker instance for generating test data
Returns:
Account: Created test account instance
"""
fake = fake or Faker()
account = Account()
account.id = fake.uuid4()
account.email = fake.email()
account.name = fake.name()
account.avatar_url = fake.url()
account.tenant_id = fake.uuid4()
account.status = "active"
account.type = "normal"
account.role = "owner"
account.interface_language = "en-US"
account.created_at = fake.date_time_this_year()
account.updated_at = account.created_at
# Create a tenant for the account
from models.account import Tenant
tenant = Tenant()
tenant.id = account.tenant_id
tenant.name = f"Test Tenant {fake.company()}"
tenant.plan = "basic"
tenant.status = "active"
tenant.created_at = fake.date_time_this_year()
tenant.updated_at = tenant.created_at
from extensions.ext_database import db
db.session.add(tenant)
db.session.add(account)
db.session.commit()
# Set the current tenant for the account
account.current_tenant = tenant
return account
def _create_test_dataset(self, db_session_with_containers, account, fake=None):
"""
Helper method to create a test dataset with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
account: The account creating the dataset
fake: Faker instance for generating test data
Returns:
Dataset: Created test dataset instance
"""
fake = fake or Faker()
dataset = Dataset()
dataset.id = fake.uuid4()
dataset.tenant_id = account.tenant_id
dataset.name = f"Test Dataset {fake.word()}"
dataset.description = fake.text(max_nb_chars=200)
dataset.provider = "vendor"
dataset.permission = "only_me"
dataset.data_source_type = "upload_file"
dataset.indexing_technique = "high_quality"
dataset.created_by = account.id
dataset.updated_by = account.id
dataset.embedding_model = "text-embedding-ada-002"
dataset.embedding_model_provider = "openai"
dataset.built_in_field_enabled = False
from extensions.ext_database import db
db.session.add(dataset)
db.session.commit()
return dataset
def _create_test_document(self, db_session_with_containers, dataset, account, fake=None):
"""
Helper method to create a test document with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
dataset: The dataset containing the document
account: The account creating the document
fake: Faker instance for generating test data
Returns:
DatasetDocument: Created test document instance
"""
fake = fake or Faker()
document = DatasetDocument()
document.id = fake.uuid4()
document.tenant_id = dataset.tenant_id
document.dataset_id = dataset.id
document.position = 1
document.data_source_type = "upload_file"
document.data_source_info = '{"upload_file_id": "test_file_id"}'
document.batch = fake.uuid4()
document.name = f"Test Document {fake.word()}.txt"
document.created_from = "upload_file"
document.created_by = account.id
document.created_api_request_id = fake.uuid4()
document.processing_started_at = fake.date_time_this_year()
document.file_id = fake.uuid4()
document.word_count = fake.random_int(min=100, max=1000)
document.parsing_completed_at = fake.date_time_this_year()
document.cleaning_completed_at = fake.date_time_this_year()
document.splitting_completed_at = fake.date_time_this_year()
document.tokens = fake.random_int(min=50, max=500)
document.indexing_started_at = fake.date_time_this_year()
document.indexing_completed_at = fake.date_time_this_year()
document.indexing_status = "completed"
document.enabled = True
document.archived = False
document.doc_form = "text_model" # Use text_model form for testing
document.doc_language = "en"
from extensions.ext_database import db
db.session.add(document)
db.session.commit()
return document
def _create_test_segments(self, db_session_with_containers, document, dataset, account, count=3, fake=None):
"""
Helper method to create test document segments with realistic data.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
document: The document containing the segments
dataset: The dataset containing the document
account: The account creating the segments
count: Number of segments to create
fake: Faker instance for generating test data
Returns:
List[DocumentSegment]: Created test segment instances
"""
fake = fake or Faker()
segments = []
for i in range(count):
segment = DocumentSegment()
segment.id = fake.uuid4()
segment.tenant_id = dataset.tenant_id
segment.dataset_id = dataset.id
segment.document_id = document.id
segment.position = i + 1
segment.content = f"Test segment content {i + 1}: {fake.text(max_nb_chars=200)}"
segment.answer = f"Test answer {i + 1}" if i % 2 == 0 else None
segment.word_count = fake.random_int(min=10, max=100)
segment.tokens = fake.random_int(min=5, max=50)
segment.keywords = [fake.word() for _ in range(3)]
segment.index_node_id = f"node_{segment.id}"
segment.index_node_hash = fake.sha256()
segment.hit_count = 0
segment.enabled = True
segment.disabled_at = None
segment.disabled_by = None
segment.status = "completed"
segment.created_by = account.id
segment.updated_by = account.id
segment.indexing_at = fake.date_time_this_year()
segment.completed_at = fake.date_time_this_year()
segment.error = None
segment.stopped_at = None
segments.append(segment)
from extensions.ext_database import db
for segment in segments:
db.session.add(segment)
db.session.commit()
return segments
def _create_dataset_process_rule(self, db_session_with_containers, dataset, fake=None):
"""
Helper method to create a dataset process rule.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
dataset: The dataset for the process rule
fake: Faker instance for generating test data
Returns:
DatasetProcessRule: Created process rule instance
"""
fake = fake or Faker()
process_rule = DatasetProcessRule()
process_rule.id = fake.uuid4()
process_rule.tenant_id = dataset.tenant_id
process_rule.dataset_id = dataset.id
process_rule.mode = "automatic"
process_rule.rules = (
"{"
'"mode": "automatic", '
'"rules": {'
'"pre_processing_rules": [], "segmentation": '
'{"separator": "\\n\\n", "max_tokens": 1000, "chunk_overlap": 50}}'
"}"
)
process_rule.created_by = dataset.created_by
process_rule.updated_by = dataset.updated_by
from extensions.ext_database import db
db.session.add(process_rule)
db.session.commit()
return process_rule
def test_disable_segments_success(self, db_session_with_containers):
"""
Test successful disabling of segments from index.
This test verifies that the task can correctly disable segments from the index
when all conditions are met, including proper index cleanup and database state updates.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_segments(db_session_with_containers, document, dataset, account, 3, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
segment_ids = [segment.id for segment in segments]
# Mock the index processor to avoid external dependencies
with patch("tasks.disable_segments_from_index_task.IndexProcessorFactory") as mock_factory:
mock_processor = MagicMock()
mock_factory.return_value.init_index_processor.return_value = mock_processor
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
mock_redis.delete.return_value = True
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Verify index processor was called correctly
mock_factory.assert_called_once_with(document.doc_form)
mock_processor.clean.assert_called_once()
# Verify the call arguments (checking by attributes rather than object identity)
call_args = mock_processor.clean.call_args
assert call_args[0][0].id == dataset.id # First argument should be the dataset
assert sorted(call_args[0][1]) == sorted(
[segment.index_node_id for segment in segments]
) # Compare sorted lists to handle any order while preserving duplicates
assert call_args[1]["with_keywords"] is True
assert call_args[1]["delete_child_chunks"] is False
# Verify Redis cache cleanup was called for each segment
assert mock_redis.delete.call_count == len(segments)
for segment in segments:
expected_key = f"segment_{segment.id}_indexing"
mock_redis.delete.assert_any_call(expected_key)
def test_disable_segments_dataset_not_found(self, db_session_with_containers):
"""
Test handling when dataset is not found.
This test ensures that the task correctly handles cases where the specified
dataset doesn't exist, logging appropriate messages and returning early.
"""
# Arrange
fake = Faker()
non_existent_dataset_id = fake.uuid4()
non_existent_document_id = fake.uuid4()
segment_ids = [fake.uuid4()]
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
# Act
result = disable_segments_from_index_task(segment_ids, non_existent_dataset_id, non_existent_document_id)
# Assert
assert result is None # Task should complete without returning a value
# Redis should not be called when dataset is not found
mock_redis.delete.assert_not_called()
def test_disable_segments_document_not_found(self, db_session_with_containers):
"""
Test handling when document is not found.
This test ensures that the task correctly handles cases where the specified
document doesn't exist, logging appropriate messages and returning early.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
non_existent_document_id = fake.uuid4()
segment_ids = [fake.uuid4()]
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, non_existent_document_id)
# Assert
assert result is None # Task should complete without returning a value
# Redis should not be called when document is not found
mock_redis.delete.assert_not_called()
def test_disable_segments_document_invalid_status(self, db_session_with_containers):
"""
Test handling when document has invalid status for disabling.
This test ensures that the task correctly handles cases where the document
is not enabled, archived, or not completed, preventing invalid operations.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_segments(db_session_with_containers, document, dataset, account, 2, fake)
# Test case 1: Document not enabled
document.enabled = False
from extensions.ext_database import db
db.session.commit()
segment_ids = [segment.id for segment in segments]
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Redis should not be called when document status is invalid
mock_redis.delete.assert_not_called()
# Test case 2: Document archived
document.enabled = True
document.archived = True
db.session.commit()
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
mock_redis.delete.assert_not_called()
# Test case 3: Document indexing not completed
document.enabled = True
document.archived = False
document.indexing_status = "indexing"
db.session.commit()
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
mock_redis.delete.assert_not_called()
def test_disable_segments_no_segments_found(self, db_session_with_containers):
"""
Test handling when no segments are found for the given IDs.
This test ensures that the task correctly handles cases where the specified
segment IDs don't exist or don't match the dataset/document criteria.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
# Use non-existent segment IDs
non_existent_segment_ids = [fake.uuid4() for _ in range(3)]
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
# Act
result = disable_segments_from_index_task(non_existent_segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Redis should not be called when no segments are found
mock_redis.delete.assert_not_called()
def test_disable_segments_index_processor_error(self, db_session_with_containers):
"""
Test handling when index processor encounters an error.
This test verifies that the task correctly handles index processor errors
by rolling back segment states and ensuring proper cleanup.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_segments(db_session_with_containers, document, dataset, account, 2, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
segment_ids = [segment.id for segment in segments]
# Mock the index processor to raise an exception
with patch("tasks.disable_segments_from_index_task.IndexProcessorFactory") as mock_factory:
mock_processor = MagicMock()
mock_processor.clean.side_effect = Exception("Index processor error")
mock_factory.return_value.init_index_processor.return_value = mock_processor
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
mock_redis.delete.return_value = True
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Verify segments were rolled back to enabled state
from extensions.ext_database import db
db.session.refresh(segments[0])
db.session.refresh(segments[1])
# Check that segments are re-enabled after error
updated_segments = db.session.query(DocumentSegment).where(DocumentSegment.id.in_(segment_ids)).all()
for segment in updated_segments:
assert segment.enabled is True
assert segment.disabled_at is None
assert segment.disabled_by is None
# Verify Redis cache cleanup was still called
assert mock_redis.delete.call_count == len(segments)
def test_disable_segments_with_different_doc_forms(self, db_session_with_containers):
"""
Test disabling segments with different document forms.
This test verifies that the task correctly handles different document forms
(paragraph, qa, parent_child) and initializes the appropriate index processor.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_segments(db_session_with_containers, document, dataset, account, 2, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
segment_ids = [segment.id for segment in segments]
# Test different document forms
doc_forms = ["text_model", "qa_model", "hierarchical_model"]
for doc_form in doc_forms:
# Update document form
document.doc_form = doc_form
from extensions.ext_database import db
db.session.commit()
# Mock the index processor factory
with patch("tasks.disable_segments_from_index_task.IndexProcessorFactory") as mock_factory:
mock_processor = MagicMock()
mock_factory.return_value.init_index_processor.return_value = mock_processor
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
mock_redis.delete.return_value = True
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
mock_factory.assert_called_with(doc_form)
def test_disable_segments_performance_timing(self, db_session_with_containers):
"""
Test that the task properly measures and logs performance timing.
This test verifies that the task correctly measures execution time
and logs performance metrics for monitoring purposes.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_segments(db_session_with_containers, document, dataset, account, 3, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
segment_ids = [segment.id for segment in segments]
# Mock the index processor
with patch("tasks.disable_segments_from_index_task.IndexProcessorFactory") as mock_factory:
mock_processor = MagicMock()
mock_factory.return_value.init_index_processor.return_value = mock_processor
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
mock_redis.delete.return_value = True
# Mock time.perf_counter to control timing
with patch("tasks.disable_segments_from_index_task.time.perf_counter") as mock_perf_counter:
mock_perf_counter.side_effect = [1000.0, 1000.5] # 0.5 seconds execution time
# Mock logger to capture log messages
with patch("tasks.disable_segments_from_index_task.logger") as mock_logger:
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Verify performance logging
mock_logger.info.assert_called()
log_calls = [call[0][0] for call in mock_logger.info.call_args_list]
performance_log = next((call for call in log_calls if "latency" in call), None)
assert performance_log is not None
assert "0.5" in performance_log # Should log the execution time
def test_disable_segments_redis_cache_cleanup(self, db_session_with_containers):
"""
Test that Redis cache is properly cleaned up for all segments.
This test verifies that the task correctly removes indexing cache entries
from Redis for all processed segments, preventing stale cache issues.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_segments(db_session_with_containers, document, dataset, account, 5, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
segment_ids = [segment.id for segment in segments]
# Mock the index processor
with patch("tasks.disable_segments_from_index_task.IndexProcessorFactory") as mock_factory:
mock_processor = MagicMock()
mock_factory.return_value.init_index_processor.return_value = mock_processor
# Mock Redis client to track delete calls
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
mock_redis.delete.return_value = True
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Verify Redis delete was called for each segment
assert mock_redis.delete.call_count == len(segments)
# Verify correct cache keys were used
expected_keys = [f"segment_{segment.id}_indexing" for segment in segments]
actual_calls = [call[0][0] for call in mock_redis.delete.call_args_list]
for expected_key in expected_keys:
assert expected_key in actual_calls
def test_disable_segments_database_session_cleanup(self, db_session_with_containers):
"""
Test that database session is properly closed after task execution.
This test verifies that the task correctly manages database sessions
and ensures proper cleanup to prevent connection leaks.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_segments(db_session_with_containers, document, dataset, account, 2, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
segment_ids = [segment.id for segment in segments]
# Mock the index processor
with patch("tasks.disable_segments_from_index_task.IndexProcessorFactory") as mock_factory:
mock_processor = MagicMock()
mock_factory.return_value.init_index_processor.return_value = mock_processor
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
mock_redis.delete.return_value = True
# Mock db.session.close to verify it's called
with patch("tasks.disable_segments_from_index_task.db.session.close") as mock_close:
# Act
result = disable_segments_from_index_task(segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Verify session was closed
mock_close.assert_called()
def test_disable_segments_empty_segment_ids(self, db_session_with_containers):
"""
Test handling when empty segment IDs list is provided.
This test ensures that the task correctly handles edge cases where
an empty list of segment IDs is provided.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
empty_segment_ids = []
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
# Act
result = disable_segments_from_index_task(empty_segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Redis should not be called when no segments are provided
mock_redis.delete.assert_not_called()
def test_disable_segments_mixed_valid_invalid_ids(self, db_session_with_containers):
"""
Test handling when some segment IDs are valid and others are invalid.
This test verifies that the task correctly processes only the valid
segment IDs and ignores invalid ones.
"""
# Arrange
fake = Faker()
account = self._create_test_account(db_session_with_containers, fake)
dataset = self._create_test_dataset(db_session_with_containers, account, fake)
document = self._create_test_document(db_session_with_containers, dataset, account, fake)
segments = self._create_test_segments(db_session_with_containers, document, dataset, account, 2, fake)
self._create_dataset_process_rule(db_session_with_containers, dataset, fake)
# Mix valid and invalid segment IDs
valid_segment_ids = [segment.id for segment in segments]
invalid_segment_ids = [fake.uuid4() for _ in range(2)]
mixed_segment_ids = valid_segment_ids + invalid_segment_ids
# Mock the index processor
with patch("tasks.disable_segments_from_index_task.IndexProcessorFactory") as mock_factory:
mock_processor = MagicMock()
mock_factory.return_value.init_index_processor.return_value = mock_processor
# Mock Redis client
with patch("tasks.disable_segments_from_index_task.redis_client") as mock_redis:
mock_redis.delete.return_value = True
# Act
result = disable_segments_from_index_task(mixed_segment_ids, dataset.id, document.id)
# Assert
assert result is None # Task should complete without returning a value
# Verify index processor was called with only valid segment node IDs
expected_node_ids = [segment.index_node_id for segment in segments]
mock_processor.clean.assert_called_once()
# Verify the call arguments
call_args = mock_processor.clean.call_args
assert call_args[0][0].id == dataset.id # First argument should be the dataset
assert sorted(call_args[0][1]) == sorted(
expected_node_ids
) # Compare sorted lists to handle any order while preserving duplicates
assert call_args[1]["with_keywords"] is True
assert call_args[1]["delete_child_chunks"] is False
# Verify Redis cleanup was called only for valid segments
assert mock_redis.delete.call_count == len(segments)

554
api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py Normal file
View File

@ -0,0 +1,554 @@
from unittest.mock import MagicMock, patch
import pytest
from faker import Faker
from extensions.ext_database import db
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, Document
from tasks.document_indexing_task import document_indexing_task
class TestDocumentIndexingTask:
"""Integration tests for document_indexing_task using testcontainers."""
@pytest.fixture
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,
):
# Setup mock indexing runner
mock_runner_instance = MagicMock()
mock_indexing_runner.return_value = mock_runner_instance
# Setup mock feature service
mock_features = MagicMock()
mock_features.billing.enabled = False
mock_feature_service.get_features.return_value = mock_features
yield {
"indexing_runner": mock_indexing_runner,
"indexing_runner_instance": mock_runner_instance,
"feature_service": mock_feature_service,
"features": mock_features,
}
def _create_test_dataset_and_documents(
self, db_session_with_containers, mock_external_service_dependencies, document_count=3
):
"""
Helper method to create a test dataset and documents for testing.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
mock_external_service_dependencies: Mock dependencies
document_count: Number of documents to create
Returns:
tuple: (dataset, documents) - Created dataset and document instances
"""
fake = Faker()
# Create account and tenant
account = Account(
email=fake.email(),
name=fake.name(),
interface_language="en-US",
status="active",
)
db.session.add(account)
db.session.commit()
tenant = Tenant(
name=fake.company(),
status="normal",
)
db.session.add(tenant)
db.session.commit()
# Create tenant-account join
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER.value,
current=True,
)
db.session.add(join)
db.session.commit()
# Create dataset
dataset = Dataset(
id=fake.uuid4(),
tenant_id=tenant.id,
name=fake.company(),
description=fake.text(max_nb_chars=100),
data_source_type="upload_file",
indexing_technique="high_quality",
created_by=account.id,
)
db.session.add(dataset)
db.session.commit()
# Create documents
documents = []
for i in range(document_count):
document = Document(
id=fake.uuid4(),
tenant_id=tenant.id,
dataset_id=dataset.id,
position=i,
data_source_type="upload_file",
batch="test_batch",
name=fake.file_name(),
created_from="upload_file",
created_by=account.id,
indexing_status="waiting",
enabled=True,
)
db.session.add(document)
documents.append(document)
db.session.commit()
# Refresh dataset to ensure it's properly loaded
db.session.refresh(dataset)
return dataset, documents
def _create_test_dataset_with_billing_features(
self, db_session_with_containers, mock_external_service_dependencies, billing_enabled=True
):
"""
Helper method to create a test dataset with billing features configured.
Args:
db_session_with_containers: Database session from testcontainers infrastructure
mock_external_service_dependencies: Mock dependencies
billing_enabled: Whether billing is enabled
Returns:
tuple: (dataset, documents) - Created dataset and document instances
"""
fake = Faker()
# Create account and tenant
account = Account(
email=fake.email(),
name=fake.name(),
interface_language="en-US",
status="active",
)
db.session.add(account)
db.session.commit()
tenant = Tenant(
name=fake.company(),
status="normal",
)
db.session.add(tenant)
db.session.commit()
# Create tenant-account join
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=TenantAccountRole.OWNER.value,
current=True,
)
db.session.add(join)
db.session.commit()
# Create dataset
dataset = Dataset(
id=fake.uuid4(),
tenant_id=tenant.id,
name=fake.company(),
description=fake.text(max_nb_chars=100),
data_source_type="upload_file",
indexing_technique="high_quality",
created_by=account.id,
)
db.session.add(dataset)
db.session.commit()
# Create documents
documents = []
for i in range(3):
document = Document(
id=fake.uuid4(),
tenant_id=tenant.id,
dataset_id=dataset.id,
position=i,
data_source_type="upload_file",
batch="test_batch",
name=fake.file_name(),
created_from="upload_file",
created_by=account.id,
indexing_status="waiting",
enabled=True,
)
db.session.add(document)
documents.append(document)
db.session.commit()
# Configure billing features
mock_external_service_dependencies["features"].billing.enabled = billing_enabled
if billing_enabled:
mock_external_service_dependencies["features"].billing.subscription.plan = "sandbox"
mock_external_service_dependencies["features"].vector_space.limit = 100
mock_external_service_dependencies["features"].vector_space.size = 50
# Refresh dataset to ensure it's properly loaded
db.session.refresh(dataset)
return dataset, documents
def test_document_indexing_task_success(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test successful document indexing with multiple documents.
This test verifies:
- Proper dataset retrieval from database
- Correct document processing and status updates
- IndexingRunner integration
- Database state updates
"""
# Arrange: Create test data
dataset, documents = self._create_test_dataset_and_documents(
db_session_with_containers, mock_external_service_dependencies, document_count=3
)
document_ids = [doc.id for doc in documents]
# Act: Execute the task
document_indexing_task(dataset.id, document_ids)
# Assert: Verify the expected outcomes
# Verify indexing runner was called correctly
mock_external_service_dependencies["indexing_runner"].assert_called_once()
mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
# Verify documents were updated to parsing status
for document in documents:
db.session.refresh(document)
assert document.indexing_status == "parsing"
assert document.processing_started_at is not None
# Verify the run method was called with correct documents
call_args = mock_external_service_dependencies["indexing_runner_instance"].run.call_args
assert call_args is not None
processed_documents = call_args[0][0] # First argument should be documents list
assert len(processed_documents) == 3
def test_document_indexing_task_dataset_not_found(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test handling of non-existent dataset.
This test verifies:
- Proper error handling for missing datasets
- Early return without processing
- Database session cleanup
- No unnecessary indexing runner calls
"""
# Arrange: Use non-existent dataset ID
fake = Faker()
non_existent_dataset_id = fake.uuid4()
document_ids = [fake.uuid4() for _ in range(3)]
# Act: Execute the task with non-existent dataset
document_indexing_task(non_existent_dataset_id, document_ids)
# Assert: Verify no processing occurred
mock_external_service_dependencies["indexing_runner"].assert_not_called()
mock_external_service_dependencies["indexing_runner_instance"].run.assert_not_called()
def test_document_indexing_task_document_not_found_in_dataset(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test handling when some documents don't exist in the dataset.
This test verifies:
- Only existing documents are processed
- Non-existent documents are ignored
- Indexing runner receives only valid documents
- Database state updates correctly
"""
# Arrange: Create test data
dataset, documents = self._create_test_dataset_and_documents(
db_session_with_containers, mock_external_service_dependencies, document_count=2
)
# Mix existing and non-existent document IDs
fake = Faker()
existing_document_ids = [doc.id for doc in documents]
non_existent_document_ids = [fake.uuid4() for _ in range(2)]
all_document_ids = existing_document_ids + non_existent_document_ids
# Act: Execute the task with mixed document IDs
document_indexing_task(dataset.id, all_document_ids)
# Assert: Verify only existing documents were processed
mock_external_service_dependencies["indexing_runner"].assert_called_once()
mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
# Verify only existing documents were updated
for document in documents:
db.session.refresh(document)
assert document.indexing_status == "parsing"
assert document.processing_started_at is not None
# Verify the run method was called with only existing documents
call_args = mock_external_service_dependencies["indexing_runner_instance"].run.call_args
assert call_args is not None
processed_documents = call_args[0][0] # First argument should be documents list
assert len(processed_documents) == 2 # Only existing documents
def test_document_indexing_task_indexing_runner_exception(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test handling of IndexingRunner exceptions.
This test verifies:
- Exceptions from IndexingRunner are properly caught
- Task completes without raising exceptions
- Database session is properly closed
- Error logging occurs
"""
# Arrange: Create test data
dataset, documents = self._create_test_dataset_and_documents(
db_session_with_containers, mock_external_service_dependencies, document_count=2
)
document_ids = [doc.id for doc in documents]
# Mock IndexingRunner to raise an exception
mock_external_service_dependencies["indexing_runner_instance"].run.side_effect = Exception(
"Indexing runner failed"
)
# Act: Execute the task
document_indexing_task(dataset.id, document_ids)
# Assert: Verify exception was handled gracefully
# The task should complete without raising exceptions
mock_external_service_dependencies["indexing_runner"].assert_called_once()
mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
# Verify documents were still updated to parsing status before the exception
for document in documents:
db.session.refresh(document)
assert document.indexing_status == "parsing"
assert document.processing_started_at is not None
def test_document_indexing_task_mixed_document_states(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test processing documents with mixed initial states.
This test verifies:
- Documents with different initial states are handled correctly
- Only valid documents are processed
- Database state updates are consistent
- IndexingRunner receives correct documents
"""
# Arrange: Create test data
dataset, base_documents = self._create_test_dataset_and_documents(
db_session_with_containers, mock_external_service_dependencies, document_count=2
)
# Create additional documents with different states
fake = Faker()
extra_documents = []
# Document with different indexing status
doc1 = Document(
id=fake.uuid4(),
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
position=2,
data_source_type="upload_file",
batch="test_batch",
name=fake.file_name(),
created_from="upload_file",
created_by=dataset.created_by,
indexing_status="completed", # Already completed
enabled=True,
)
db.session.add(doc1)
extra_documents.append(doc1)
# Document with disabled status
doc2 = Document(
id=fake.uuid4(),
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
position=3,
data_source_type="upload_file",
batch="test_batch",
name=fake.file_name(),
created_from="upload_file",
created_by=dataset.created_by,
indexing_status="waiting",
enabled=False, # Disabled
)
db.session.add(doc2)
extra_documents.append(doc2)
db.session.commit()
all_documents = base_documents + extra_documents
document_ids = [doc.id for doc in all_documents]
# Act: Execute the task with mixed document states
document_indexing_task(dataset.id, document_ids)
# Assert: Verify processing
mock_external_service_dependencies["indexing_runner"].assert_called_once()
mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
# Verify all documents were updated to parsing status
for document in all_documents:
db.session.refresh(document)
assert document.indexing_status == "parsing"
assert document.processing_started_at is not None
# Verify the run method was called with all documents
call_args = mock_external_service_dependencies["indexing_runner_instance"].run.call_args
assert call_args is not None
processed_documents = call_args[0][0] # First argument should be documents list
assert len(processed_documents) == 4
def test_document_indexing_task_billing_sandbox_plan_batch_limit(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test billing validation for sandbox plan batch upload limit.
This test verifies:
- Sandbox plan batch upload limit enforcement
- Error handling for batch upload limit exceeded
- Document status updates to error state
- Proper error message recording
"""
# Arrange: Create test data with billing enabled
dataset, documents = self._create_test_dataset_with_billing_features(
db_session_with_containers, mock_external_service_dependencies, billing_enabled=True
)
# Configure sandbox plan with batch limit
mock_external_service_dependencies["features"].billing.subscription.plan = "sandbox"
# Create more documents than sandbox plan allows (limit is 1)
fake = Faker()
extra_documents = []
for i in range(2): # Total will be 5 documents (3 existing + 2 new)
document = Document(
id=fake.uuid4(),
tenant_id=dataset.tenant_id,
dataset_id=dataset.id,
position=i + 3,
data_source_type="upload_file",
batch="test_batch",
name=fake.file_name(),
created_from="upload_file",
created_by=dataset.created_by,
indexing_status="waiting",
enabled=True,
)
db.session.add(document)
extra_documents.append(document)
db.session.commit()
all_documents = documents + extra_documents
document_ids = [doc.id for doc in all_documents]
# Act: Execute the task with too many documents for sandbox plan
document_indexing_task(dataset.id, document_ids)
# Assert: Verify error handling
for document in all_documents:
db.session.refresh(document)
assert document.indexing_status == "error"
assert document.error is not None
assert "batch upload" in document.error
assert document.stopped_at is not None
# Verify no indexing runner was called
mock_external_service_dependencies["indexing_runner"].assert_not_called()
def test_document_indexing_task_billing_disabled_success(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test successful processing when billing is disabled.
This test verifies:
- Processing continues normally when billing is disabled
- No billing validation occurs
- Documents are processed successfully
- IndexingRunner is called correctly
"""
# Arrange: Create test data with billing disabled
dataset, documents = self._create_test_dataset_with_billing_features(
db_session_with_containers, mock_external_service_dependencies, billing_enabled=False
)
document_ids = [doc.id for doc in documents]
# Act: Execute the task with billing disabled
document_indexing_task(dataset.id, document_ids)
# Assert: Verify successful processing
mock_external_service_dependencies["indexing_runner"].assert_called_once()
mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
# Verify documents were updated to parsing status
for document in documents:
db.session.refresh(document)
assert document.indexing_status == "parsing"
assert document.processing_started_at is not None
def test_document_indexing_task_document_is_paused_error(
self, db_session_with_containers, mock_external_service_dependencies
):
"""
Test handling of DocumentIsPausedError from IndexingRunner.
This test verifies:
- DocumentIsPausedError is properly caught and handled
- Task completes without raising exceptions
- Appropriate logging occurs
- Database session is properly closed
"""
# Arrange: Create test data
dataset, documents = self._create_test_dataset_and_documents(
db_session_with_containers, mock_external_service_dependencies, document_count=2
)
document_ids = [doc.id for doc in documents]
# Mock IndexingRunner to raise DocumentIsPausedError
from core.indexing_runner import DocumentIsPausedError
mock_external_service_dependencies["indexing_runner_instance"].run.side_effect = DocumentIsPausedError(
"Document indexing is paused"
)
# Act: Execute the task
document_indexing_task(dataset.id, document_ids)
# Assert: Verify exception was handled gracefully
# The task should complete without raising exceptions
mock_external_service_dependencies["indexing_runner"].assert_called_once()
mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once()
# Verify documents were still updated to parsing status before the exception
for document in documents:
db.session.refresh(document)
assert document.indexing_status == "parsing"
assert document.processing_started_at is not None